[pypy-commit] pypy unicode-dtype: Implement UnicodeType.fill()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: unicode-dtype
Changeset: r78019:f3f7e88f6b92
Date: 2015-06-10 05:12 +0100
http://bitbucket.org/pypy/pypy/changeset/f3f7e88f6b92/

Log:	Implement UnicodeType.fill()

diff --git a/pypy/module/micronumpy/test/test_ndarray.py b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -389,6 +389,9 @@
         assert zeros((), dtype='S') == ''
         assert zeros((), dtype='S').shape == ()
         assert zeros((), dtype='S').dtype == '|S1'
+        assert zeros(5, dtype='U')[4] == u''
+        assert zeros(5, dtype='U').shape == (5,)
+        assert zeros(5, dtype='U').dtype == '<U1'
 
     def test_check_shape(self):
         import numpy as np
@@ -2432,6 +2435,12 @@
         a.fill(12)
         assert (a == '1').all()
 
+    def test_unicode_filling(self):
+        import numpy as np
+        a = np.empty((10,10), dtype='U1')
+        a.fill(12)
+        assert (a == u'1').all()
+
     def test_boolean_indexing(self):
         import numpy as np
         a = np.zeros((1, 3))
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -50,6 +50,7 @@
             pass
         return _raw_storage_getitem_unaligned(T, storage, offset)
 '''
+
 def simple_unary_op(func):
     specialize.argtype(1)(func)
     @functools.wraps(func)
@@ -2195,10 +2196,20 @@
     def store(self, arr, i, offset, box, native):
         assert isinstance(box, boxes.W_UnicodeBox)
         value = box._value
-        for k in range(len(value)):
+        with arr as storage:
+            self._store(storage, i, offset, box, arr.dtype.elsize)
+
+    @jit.unroll_safe
+    def _store(self, storage, i, offset, box, width):
+        size = min(width // 4, len(box._value))
+        for k in range(size):
             index = i + offset + 4*k
             data = rffi.cast(Int32.T, ord(box._value[k]))
-            raw_storage_setitem_unaligned(arr.storage, index, data)
+            raw_storage_setitem_unaligned(storage, index, data)
+        for k in range(size, width // 4):
+            index = i + offset + 4*k
+            data = rffi.cast(Int32.T, 0)
+            raw_storage_setitem_unaligned(storage, index, data)
 
     def read(self, arr, i, offset, dtype):
         if dtype is None:
@@ -2269,7 +2280,9 @@
         raise NotImplementedError
 
     def fill(self, storage, width, native, box, start, stop, offset, gcstruct):
-        raise NotImplementedError
+        assert isinstance(box, boxes.W_UnicodeBox)
+        for i in xrange(start, stop, width):
+            self._store(storage, i, offset, box, width)
 
 
 class VoidType(FlexibleType):