[pypy-commit] pypy default: import to_coords from numppy-flatitter branch and modify it. Thanks mattip!

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r51805:9183b721128a
Date: 2012-01-26 20:37 +0200
http://bitbucket.org/pypy/pypy/changeset/9183b721128a/

Log:	import to_coords from numppy-flatitter branch and modify it. Thanks
	mattip!

diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -1,4 +1,5 @@
 from pypy.rlib import jit
+from pypy.interpreter.error import OperationError
 
 
 @jit.look_inside_iff(lambda shape, start, strides, backstrides, chunks:
@@ -37,3 +38,27 @@
     rstrides = [0] * (len(res_shape) - len(orig_shape)) + rstrides
     rbackstrides = [0] * (len(res_shape) - len(orig_shape)) + rbackstrides
     return rstrides, rbackstrides
+
+def to_coords(space, shape, size, order, w_item_or_slice):
+    '''Returns a start coord, step, and length.
+    '''
+    start = lngth = step = 0
+    if not (space.isinstance_w(w_item_or_slice, space.w_int) or
+        space.isinstance_w(w_item_or_slice, space.w_slice)):
+        raise OperationError(space.w_IndexError,
+                             space.wrap('unsupported iterator index'))
+            
+    start, stop, step, lngth = space.decode_index4(w_item_or_slice, size)
+        
+    coords = [0] * len(shape)
+    i = start
+    if order == 'C':
+        for s in range(len(shape) -1, -1, -1):
+            coords[s] = i % shape[s]
+            i //= shape[s]
+    else:
+        for s in range(len(shape)):
+            coords[s] = i % shape[s]
+            i //= shape[s]
+
+    return coords, step, lngth
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
@@ -167,6 +167,24 @@
         assert calc_new_strides([1, 1, 105, 1, 1], [7, 15], [1, 7],'F') == \
                                     [1, 1, 1, 105, 105]
 
+    def test_to_coords(self):
+        from pypy.module.micronumpy.strides import to_coords
+
+        def _to_coords(index, order):
+            return to_coords(self.space, [2, 3, 4], 24, order,
+                             self.space.wrap(index))[0]
+        
+        assert _to_coords(0, 'C') == [0, 0, 0]
+        assert _to_coords(1, 'C') == [0, 0, 1]
+        assert _to_coords(-1, 'C') == [1, 2, 3]
+        assert _to_coords(5, 'C') == [0, 1, 1]
+        assert _to_coords(13, 'C') == [1, 0, 1]
+        assert _to_coords(0, 'F') == [0, 0, 0]
+        assert _to_coords(1, 'F') == [1, 0, 0]
+        assert _to_coords(-1, 'F') == [1, 2, 3]
+        assert _to_coords(5, 'F') == [1, 2, 0]
+        assert _to_coords(13, 'F') == [1, 0, 2]
+        
 
 class AppTestNumArray(BaseNumpyAppTest):
     def test_ndarray(self):