Author: Maciej Fijalkowski <fij...@gmail.com>
Branch:
Changeset: r54146:5da12de5d696
Date: 2012-04-02 10:41 +0200
http://bitbucket.org/pypy/pypy/changeset/5da12de5d696/
Log: implement repeat. a tad inefficient, but I could not care less
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
@@ -1261,6 +1261,31 @@
return convert_to_array(space, w_obj2).descr_dot(space, w_arr)
return w_arr.descr_dot(space, w_obj2)
+@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)
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,15 @@
assert (ones(1) + ones(1)).nbytes == 8
assert array(3.0).nbytes == 8
+ def test_repeat(self):
+ from _numpypy import repeat
+ 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()
+
class AppTestMultiDim(BaseNumpyAppTest):
def test_init(self):
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
http://mail.python.org/mailman/listinfo/pypy-commit
