BUG: Update base-traversal algorithm for array sub-classes so as to stop the base-traversal when the new base would not be an instance of the sub-class.

teoliphant · certik · commit a9c7b38f293e · 2012-11-13T18:03:14.000-08:00
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
@@ -143,29 +143,38 @@ PyArray_SetBaseObject(PyArrayObject *arr, PyObject *obj)
     }
 
     /*
-     * Don't allow chains of views, always set the base
-     * to the owner of the data. That is, either the first object
-     * which isn't an array, or the first object which owns
-     * its own data.
+     * Don't allow infinite chains of views, always set the base
+     * to the first owner of the data.  
+     * That is, either the first object which isn't an array, 
+     * or the first object which owns its own data.
      */
+
     while (PyArray_Check(obj) && (PyObject *)arr != obj) {
         PyArrayObject *obj_arr = (PyArrayObject *)obj;
         PyObject *tmp;
 
         /* Propagate WARN_ON_WRITE through views. */
         if (PyArray_FLAGS(obj_arr) & NPY_ARRAY_WARN_ON_WRITE) {
             PyArray_ENABLEFLAGS(arr, NPY_ARRAY_WARN_ON_WRITE);
-        }
+        }   
 
         /* If this array owns its own data, stop collapsing */
         if (PyArray_CHKFLAGS(obj_arr, NPY_ARRAY_OWNDATA)) {
             break;
-        }
-        /* If there's no base, stop collapsing */
+        }   
+
         tmp = PyArray_BASE(obj_arr);
+        /* If there's no base, stop collapsing */
         if (tmp == NULL) {
             break;
         }
+        /* Stop the collapse for array sub-classes if new base
+         *   would not be of the same type. 
+         */
+        if (!(PyArray_CheckExact(arr)) & (Py_TYPE(tmp) != Py_TYPE(arr))) {
+            break;
+        }
+
 
         Py_INCREF(tmp);
         Py_DECREF(obj);
diff --git a/numpy/core/tests/test_memmap.py b/numpy/core/tests/test_memmap.py
@@ -104,5 +104,12 @@ def test_slicing_keeps_references(self):
                     shape=self.shape)
         assert fp[:2, :2]._mmap is fp._mmap
 
+    def test_view(self):
+        fp = memmap(self.tmpfp, dtype=self.dtype, shape=self.shape)
+        new1 = fp.view()
+        new2 = new1.view()
+        assert(new1.base is fp)
+        assert(new2.base is fp)
+
 if __name__ == "__main__":
     run_module_suite()
