[pypy-commit] pypy default: test, fix for matrix subtype which modfies shape in __getitem__

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: 
Changeset: r68441:bacd9679adbf
Date: 2013-12-15 07:43 +0200
http://bitbucket.org/pypy/pypy/changeset/bacd9679adbf/

Log:	test, fix for matrix subtype which modfies shape in __getitem__
	causing an infinite loop in find_shape_and_elems

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
@@ -51,8 +51,8 @@
             rstrides.append(strides[i])
             rbackstrides.append(backstrides[i])
     if backwards:
-        rstrides = rstrides + [0] * (len(res_shape) - len(orig_shape))  
-        rbackstrides = rbackstrides + [0] * (len(res_shape) - len(orig_shape)) 
+        rstrides = rstrides + [0] * (len(res_shape) - len(orig_shape))
+        rbackstrides = rbackstrides + [0] * (len(res_shape) - len(orig_shape))
     else:
         rstrides = [0] * (len(res_shape) - len(orig_shape)) + rstrides
         rbackstrides = [0] * (len(res_shape) - len(orig_shape)) + rbackstrides
@@ -62,7 +62,7 @@
     if (is_rec_type and space.isinstance_w(w_elem, space.w_tuple)):
         return True
     if (space.isinstance_w(w_elem, space.w_tuple) or
-        isinstance(w_elem, W_NDimArray) or    
+        isinstance(w_elem, W_NDimArray) or
         space.isinstance_w(w_elem, space.w_list)):
         return False
     return True
@@ -87,6 +87,12 @@
                 space.len_w(w_elem) != size):
                 raise OperationError(space.w_ValueError, space.wrap(
                     "setting an array element with a sequence"))
+            w_array = space.lookup(w_elem, '__array__')
+            if w_array is not None:
+                # Make sure we call the array implementation of listview,
+                # since for some ndarray subclasses (matrix, for instance)
+                # listview does not reduce but rather returns the same class
+                w_elem = space.get_and_call_function(w_array, w_elem, space.w_None)
             new_batch += space.listview(w_elem)
         shape.append(size)
         batch = new_batch
diff --git a/pypy/module/micronumpy/test/test_subtype.py b/pypy/module/micronumpy/test/test_subtype.py
--- a/pypy/module/micronumpy/test/test_subtype.py
+++ b/pypy/module/micronumpy/test/test_subtype.py
@@ -258,3 +258,46 @@
             assert isinstance(b, D)
         c = array(a, float)
         assert c.dtype is dtype(float)
+
+    def test__getitem_modifies_shape(self):
+        import numpypy as N
+        # numpy's matrix class caused an infinite loop
+        class matrix(N.ndarray):
+            getcnt = 0
+            def __new__(subtype, data, dtype=None, copy=True):
+                arr = N.array(data, dtype=dtype, copy=copy)
+                shape = arr.shape
+
+                ret = N.ndarray.__new__(subtype, shape, arr.dtype,
+                                        buffer=arr,
+                                        order=True)
+                return ret
+
+            def __getitem__(self, index):
+                matrix.getcnt += 1
+                if matrix.getcnt > 10:
+                    # XXX strides.find_shape_and_elems is sensitive
+                    # to shape modification
+                    xxx
+                out = N.ndarray.__getitem__(self, index)
+
+                if not isinstance(out, N.ndarray):
+                    return out
+                    # Determine when we should have a column array
+                old_shape = out.shape
+                if out.ndim < 2:
+                    sh = out.shape[0]
+                    try:
+                        n = len(index)
+                    except:
+                        n = 0
+                    if n > 1:
+                        out.shape = (sh, 1)
+                    else:
+                        out.shape = (1, sh)
+                print 'out, shape was',old_shape,'now',out.shape
+                return out
+        a = matrix([[1., 2.]])
+        b = N.array([a])
+
+