[pypy-commit] pypy ndarray-subtype: a branch to allow subtype of ndarray, add tests and start to implement

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: ndarray-subtype
Changeset: r65174:38c0e3d688bc
Date: 2013-07-03 18:31 +0300
http://bitbucket.org/pypy/pypy/changeset/38c0e3d688bc/

Log:	a branch to allow subtype of ndarray, add tests and start to
	implement

diff --git a/pypy/module/micronumpy/base.py b/pypy/module/micronumpy/base.py
--- a/pypy/module/micronumpy/base.py
+++ b/pypy/module/micronumpy/base.py
@@ -23,7 +23,7 @@
         self.implementation = implementation
 
     @staticmethod
-    def from_shape(shape, dtype, order='C'):
+    def from_shape(shape, dtype, order='C', subtype_and_space=(None, None)):
         from pypy.module.micronumpy.arrayimpl import concrete, scalar
 
         if not shape:
@@ -32,10 +32,16 @@
             strides, backstrides = calc_strides(shape, dtype.base, order)
             impl = concrete.ConcreteArray(shape, dtype.base, order, strides,
                                       backstrides)
+        if subtype_and_space[0]:
+            space = subtype_and_space[1]
+            subtype = subtype_and_space[0]
+            ret = space.allocate_instance(W_NDimArray, subtype)
+            W_NDimArray.__init__(ret, impl)
+            return ret
         return W_NDimArray(impl)
 
     @staticmethod
-    def from_shape_and_storage(shape, storage, dtype, order='C', owning=False):
+    def from_shape_and_storage(shape, storage, dtype, order='C', owning=False, subtype_and_space=(None, None)):
         from pypy.module.micronumpy.arrayimpl import concrete
         assert shape
         strides, backstrides = calc_strides(shape, dtype, order)
@@ -46,6 +52,13 @@
         else:
             impl = concrete.ConcreteArrayNotOwning(shape, dtype, order, strides,
                                                 backstrides, storage)
+        if subtype_and_space[0]:
+            print 'creating subclass',subtype_and_space
+            space = subtype_and_space[1]
+            subtype = subtype_and_space[0]
+            ret = space.allocate_instance(W_NDimArray, subtype)
+            W_NDimArray.__init__(ret, impl)
+            return ret
         return W_NDimArray(impl)
 
     @staticmethod
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
@@ -914,7 +914,7 @@
     dtype = space.interp_w(interp_dtype.W_Dtype,
                            space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype))
     shape = _find_shape(space, w_shape, dtype)
-    return W_NDimArray.from_shape_and_storage(shape, storage, dtype)
+    return W_NDimArray.from_shape_and_storage(shape, storage, dtype, (space, w_cls))
 
 W_NDimArray.typedef = TypeDef(
     "ndarray",
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
@@ -1442,7 +1442,7 @@
         assert x.view('int8').shape == (10, 3)
 
     def test_ndarray_view_empty(self):
-        from numpypy import array, int8, int16, dtype
+        from numpypy import array, int8, int16
         x = array([], dtype=[('a', int8), ('b', int8)])
         y = x.view(dtype=int16)
 
@@ -1458,6 +1458,20 @@
         skip('not implemented yet')
         assert s.view('double') < 7e-323
 
+    def test_subclass_view(self):
+        from numpypy import ndarray, array
+        class matrix(ndarray):
+            def __new__(subtype, data, dtype=None, copy=True):
+                print 'matix.__new__(',subtype,',',data,'...)'
+                if isinstance(data, matrix):
+                    return data
+                return data.view(subtype)
+        a = array(range(5))
+        b = matrix(a)
+        print type(b),b
+        assert False
+        assert (b == a).all()
+
     def test_tolist_scalar(self):
         from numpypy import int32, bool_
         x = int32(23)
@@ -2871,6 +2885,12 @@
         assert y[0, 1] == 2
         y[0, 1] = 42
         assert x[1] == 42
+        class C(ndarray):
+            pass
+        z = ndarray._from_shape_and_storage([4, 1], addr, x.dtype, C)
+        assert isinstance(z, C)
+        assert z.shape == (4, 1)
+        assert z[1, 0] == 42
 
     def test___pypy_data__(self):
         from numpypy import array