[pypy-commit] pypy numpy-subarrays: Implement dtype.shape

[rguillebert]

Author: Romain Guillebert <romain.py at gmail.com>
Branch: numpy-subarrays
Changeset: r63986:159e1024713b
Date: 2013-05-11 18:08 +0200
http://bitbucket.org/pypy/pypy/changeset/159e1024713b/

Log:	Implement dtype.shape

diff --git a/pypy/module/micronumpy/interp_dtype.py b/pypy/module/micronumpy/interp_dtype.py
--- a/pypy/module/micronumpy/interp_dtype.py
+++ b/pypy/module/micronumpy/interp_dtype.py
@@ -50,7 +50,7 @@
 
     def __init__(self, itemtype, num, kind, name, char, w_box_type,
                  alternate_constructors=[], aliases=[],
-                 fields=None, fieldnames=None, native=True):
+                 fields=None, fieldnames=None, native=True, shape=[], subdtype=None):
         self.itemtype = itemtype
         self.num = num
         self.kind = kind
@@ -63,6 +63,8 @@
         self.fieldnames = fieldnames
         self.native = native
         self.float_type = None
+        self.shape = shape
+        self.subdtype = subdtype
 
     @specialize.argtype(1)
     def box(self, value):
@@ -111,8 +113,11 @@
     def descr_get_alignment(self, space):
         return space.wrap(self.itemtype.alignment)
 
+    def descr_get_subdtype(self, space):
+        return space.newtuple([space.wrap(self.subdtype), space.newtuple(shape)])
+
     def descr_get_shape(self, space):
-        return space.newtuple([])
+        return space.newtuple(self.shape)
 
     def eq(self, space, w_other):
         w_other = space.call_function(space.gettypefor(W_Dtype), w_other)
@@ -279,17 +284,23 @@
     ofs_and_items = []
     fieldnames = []
     for w_elem in lst_w:
-        w_fldname, w_flddesc = space.fixedview(w_elem, 2)
-        subdtype = descr__new__(space, space.gettypefor(W_Dtype), w_flddesc)
+        size = 1
+        shape = space.newtuple([])
+        if space.len_w(w_elem) == 3:
+            w_shape = space.getitem(w_elem, space.wrap(2))
+        w_fldname = space.getitem(w_elem, space.wrap(0))
+        w_flddesc = space.getitem(w_elem, space.wrap(1))
+        subdtype = descr__new__(space, space.gettypefor(W_Dtype), w_flddesc, w_shape=w_shape)
         fldname = space.str_w(w_fldname)
         if fldname in fields:
             raise OperationError(space.w_ValueError, space.wrap("two fields with the same name"))
         assert isinstance(subdtype, W_Dtype)
         fields[fldname] = (offset, subdtype)
         ofs_and_items.append((offset, subdtype.itemtype))
-        offset += subdtype.itemtype.get_element_size()
+        offset += subdtype.itemtype.get_element_size() * size
         fieldnames.append(fldname)
     itemtype = types.RecordType(ofs_and_items, offset)
+    import pdb; pdb.set_trace()
     return W_Dtype(itemtype, 20, VOIDLTR, "void" + str(8 * itemtype.get_element_size()),
                    "V", space.gettypefor(interp_boxes.W_VoidBox), fields=fields,
                    fieldnames=fieldnames)
@@ -333,10 +344,19 @@
         raise OperationError(space.w_NotImplementedError, space.wrap(
             "dtype from spec"))
 
-def descr__new__(space, w_subtype, w_dtype, w_align=None, w_copy=None):
+def descr__new__(space, w_subtype, w_dtype, w_align=None, w_copy=None, w_shape=None):
     # w_align and w_copy are necessary for pickling
     cache = get_dtype_cache(space)
 
+    if w_shape is not None and space.len_w(w_shape) > 0:
+        subdtype = descr__new__(space, w_subtype, w_dtype, w_align, w_copy)
+        size = 1
+        shape = space.listview(w_shape)
+        for dim in shape:
+            size *= space.int_w(dim)
+        return W_Dtype(types.VoidType(subdtype.itemtype.get_element_size() * size), 20, VOIDLTR, "void" + str(8 * subdtype.itemtype.get_element_size() * size),
+                    "V", space.gettypefor(interp_boxes.W_VoidBox), shape=shape, subdtype=subdtype)
+
     if space.is_none(w_dtype):
         return cache.w_float64dtype
     elif space.isinstance_w(w_dtype, w_subtype):
@@ -391,6 +411,7 @@
     name = interp_attrproperty('name', cls=W_Dtype),
     fields = GetSetProperty(W_Dtype.descr_get_fields),
     names = GetSetProperty(W_Dtype.descr_get_names),
+    subdtype = GetSetProperty(W_Dtype.descr_get_subdtype),
 )
 W_Dtype.typedef.acceptable_as_base_class = False
 
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -1076,7 +1076,7 @@
 
     def to_builtin_type(self, space, box):
         real,imag = self.for_computation(self.unbox(box))
-        return space.newcomplex(real, imag) 
+        return space.newcomplex(real, imag)
 
     def read_bool(self, arr, i, offset):
         v = self.for_computation(self._read(arr.storage, i, offset))
@@ -1217,7 +1217,7 @@
 
     @raw_binary_op
     def le(self, v1, v2):
-        return self._lt(v1, v2) or self._eq(v1, v2) 
+        return self._lt(v1, v2) or self._eq(v1, v2)
 
     @raw_binary_op
     def gt(self, v1, v2):
@@ -1225,7 +1225,7 @@
 
     @raw_binary_op
     def ge(self, v1, v2):
-        return self._lt(v2, v1) or self._eq(v2, v1) 
+        return self._lt(v2, v1) or self._eq(v2, v1)
 
     def _bool(self, v):
         return bool(v[0]) or bool(v[1])
@@ -1341,7 +1341,7 @@
             return rcomplex.c_div((v[0], -v[1]), (a2, 0.))
         except ZeroDivisionError:
             return rfloat.NAN, rfloat.NAN
- 
+
     # No floor, ceil, trunc in numpy for complex
     #@simple_unary_op
     #def floor(self, v):
@@ -1696,7 +1696,7 @@
         for j in range(i + 1, self.size):
             arr.storage[j] = '\x00'
         return interp_boxes.W_StringBox(arr,  0, arr.dtype)
-        
+
 class VoidType(BaseType, BaseStringType):
     T = lltype.Char