[pypy-commit] pypy default: merge

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r54148:5708f4c44e53
Date: 2012-04-02 12:10 +0200
http://bitbucket.org/pypy/pypy/changeset/5708f4c44e53/

Log:	merge

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -29,6 +29,7 @@
         'flatiter': 'interp_numarray.W_FlatIterator',
         'isna': 'interp_numarray.isna',
         'concatenate': 'interp_numarray.concatenate',
+        'repeat': 'interp_numarray.repeat',
 
         'set_string_function': 'appbridge.set_string_function',
 
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -673,6 +673,10 @@
     def compute_first_step(self, sig, frame):
         pass
 
+    @unwrap_spec(repeats=int)
+    def descr_repeat(self, space, repeats, w_axis=None):
+        return repeat(space, self, repeats, w_axis)
+
 def convert_to_array(space, w_obj):
     if isinstance(w_obj, BaseArray):
         return w_obj
@@ -1261,6 +1265,31 @@
         return convert_to_array(space, w_obj2).descr_dot(space, w_arr)
     return w_arr.descr_dot(space, w_obj2)
 
+ at unwrap_spec(repeats=int)
+def repeat(space, w_arr, repeats, w_axis=None):
+    arr = convert_to_array(space, w_arr)
+    if space.is_w(w_axis, space.w_None):
+        arr = arr.descr_flatten(space).get_concrete()
+        orig_size = arr.shape[0]
+        shape = [arr.shape[0] * repeats]
+        res = W_NDimArray(shape, arr.find_dtype())
+        for i in range(repeats):
+            Chunks([Chunk(i, shape[0] - repeats + i, repeats,
+                          orig_size)]).apply(res).setslice(space, arr)
+    else:
+        arr = arr.get_concrete()
+        axis = space.int_w(w_axis)
+        shape = arr.shape[:]
+        chunks = [Chunk(0, i, 1, i) for i in shape]
+        orig_size = shape[axis]
+        shape[axis] *= repeats
+        res = W_NDimArray(shape, arr.find_dtype())
+        for i in range(repeats):
+            chunks[axis] = Chunk(i, shape[axis] - repeats + i, repeats,
+                                 orig_size)
+            Chunks(chunks).apply(res).setslice(space, arr)
+    return res
+
 @unwrap_spec(axis=int)
 def concatenate(space, w_args, axis=0):
     args_w = space.listview(w_args)
@@ -1386,6 +1415,7 @@
     tolist = interp2app(BaseArray.descr_tolist),
     take = interp2app(BaseArray.descr_take),
     compress = interp2app(BaseArray.descr_compress),
+    repeat = interp2app(BaseArray.descr_repeat),
 )
 
 
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -1399,6 +1399,16 @@
         assert (ones(1) + ones(1)).nbytes == 8
         assert array(3.0).nbytes == 8
 
+    def test_repeat(self):
+        from _numpypy import repeat, array
+        assert (repeat([[1, 2], [3, 4]], 3) == [1, 1, 1, 2, 2, 2,
+                                                3, 3, 3, 4, 4, 4]).all()
+        assert (repeat([[1, 2], [3, 4]], 2, axis=0) == [[1, 2], [1, 2], [3, 4],
+                                                        [3, 4]]).all()
+        assert (repeat([[1, 2], [3, 4]], 2, axis=1) == [[1, 1, 2, 2], [3, 3,
+                                                        4, 4]]).all()
+        assert (array([1, 2]).repeat(2) == array([1, 1, 2, 2])).all()
+
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):