[pypy-commit] pypy numpy-refactor: repeat;

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-refactor
Changeset: r57164:221a7a7f1ed5
Date: 2012-09-05 22:15 +0200
http://bitbucket.org/pypy/pypy/changeset/221a7a7f1ed5/

Log:	repeat;

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
@@ -28,15 +28,12 @@
         'dot': 'interp_arrayops.dot',
         'fromstring': 'interp_support.fromstring',
         'flatiter': 'interp_flatiter.W_FlatIterator',
-        'isna': 'interp_numarray.isna',
         'concatenate': 'interp_arrayops.concatenate',
-        'repeat': 'interp_numarray.repeat',
+        'repeat': 'interp_arrayops.repeat',
         'where': 'interp_arrayops.where',
 
         'set_string_function': 'appbridge.set_string_function',
 
-        'count_reduce_items': 'interp_numarray.count_reduce_items',
-
         'True_': 'types.Bool.True',
         'False_': 'types.Bool.False',
 
diff --git a/pypy/module/micronumpy/interp_arrayops.py b/pypy/module/micronumpy/interp_arrayops.py
--- a/pypy/module/micronumpy/interp_arrayops.py
+++ b/pypy/module/micronumpy/interp_arrayops.py
@@ -112,3 +112,27 @@
         Chunks(chunks).apply(res.implementation).implementation.setslice(space, arr)
         axis_start += arr.get_shape()[axis]
     return res
+
+ 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)
+        orig_size = arr.get_shape()[0]
+        shape = [arr.get_shape()[0] * repeats]
+        res = W_NDimArray.from_shape(shape, arr.get_dtype())
+        for i in range(repeats):
+            Chunks([Chunk(i, shape[0] - repeats + i, repeats,
+                          orig_size)]).apply(res.implementation).implementation.setslice(space, arr)
+    else:
+        axis = space.int_w(w_axis)
+        shape = arr.get_shape()[:]
+        chunks = [Chunk(0, i, 1, i) for i in shape]
+        orig_size = shape[axis]
+        shape[axis] *= repeats
+        res = W_NDimArray.from_shape(shape, arr.get_dtype())
+        for i in range(repeats):
+            chunks[axis] = Chunk(i, shape[axis] - repeats + i, repeats,
+                                 orig_size)
+            Chunks(chunks).apply(res.implementation).implementation.setslice(space, arr)
+    return res
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
@@ -10,6 +10,7 @@
 from pypy.module.micronumpy.appbridge import get_appbridge_cache
 from pypy.module.micronumpy import loop
 from pypy.module.micronumpy.dot import match_dot_shapes
+from pypy.module.micronumpy.interp_arrayops import repeat
 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib import jit
 from pypy.rlib.rstring import StringBuilder
@@ -185,6 +186,10 @@
             return w_res.descr_copy(space)
         return w_res
 
+    @unwrap_spec(repeats=int)
+    def descr_repeat(self, space, repeats, w_axis=None):
+        return repeat(space, self, repeats, w_axis)
+
     # --------------------- operations ----------------------------
 
     def _unaryop_impl(ufunc_name):
@@ -433,6 +438,7 @@
     tolist = interp2app(W_NDimArray.descr_tolist),
     flatten = interp2app(W_NDimArray.descr_flatten),
     ravel = interp2app(W_NDimArray.descr_ravel),
+    repeat = interp2app(W_NDimArray.descr_repeat),
 )
 
 @unwrap_spec(ndmin=int, copy=bool, subok=bool)
@@ -501,11 +507,3 @@
     arr.fill(one)
     return space.wrap(arr)
 
-def isna(space):
-    pass
-
-def repeat(space):
-    pass
-
-def count_reduce_items(space):
-    pass