[pypy-commit] pypy matrixmath: make 0-dimension arrays (scalars) numpy-compatible

[mattip]

Author: mattip
Branch: matrixmath
Changeset: r49864:5ffa83378d04
Date: 2011-11-27 22:32 +0200
http://bitbucket.org/pypy/pypy/changeset/5ffa83378d04/

Log:	make 0-dimension arrays (scalars) numpy-compatible

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
@@ -227,7 +227,7 @@
         self.strides = []
         self.backstrides = []
         for i in range(len(arr.shape)):
-            if arr.shape[i]==1:
+            if arr.shape[i] == 1:
                 self.strides.append(0)
                 self.backstrides.append(0)
             else:
@@ -683,6 +683,9 @@
     def descr_getitem(self, space, w_idx):
         if self._single_item_result(space, w_idx):
             concrete = self.get_concrete()
+            if len(concrete.shape) < 1:
+                raise OperationError(space.w_IndexError, space.wrap(
+                        "0-d arrays can't be indexed"))
             item = concrete._index_of_single_item(space, w_idx)
             return concrete.getitem(item).wrap(space)
         chunks = self._prepare_slice_args(space, w_idx)
@@ -692,6 +695,9 @@
         self.invalidated()
         concrete = self.get_concrete()
         if self._single_item_result(space, w_idx):
+            if len(concrete.shape) < 1:
+                raise OperationError(space.w_IndexError, space.wrap(
+                        "0-d arrays can't be indexed"))
             item = concrete._index_of_single_item(space, w_idx)
             concrete.setitem_w(space, item, w_value)
             return
@@ -759,17 +765,20 @@
             self.start_iter(self.shape)).wrap(space)))
 
     def descr_get_transpose(self, space):
+        concrete = self.get_concrete()
+        if len(concrete.shape) < 2:
+            return space.wrap(self)
         new_sig = signature.Signature.find_sig([
             NDimSlice.signature, self.signature
         ])
         strides = []
         backstrides = []
         shape = []
-        for i in range(len(self.shape) - 1, -1, -1):
-            strides.append(self.strides[i])
-            backstrides.append(self.backstrides[i])
-            shape.append(self.shape[i])
-        return space.wrap(NDimSlice(self, new_sig, self.start, strides[:], 
+        for i in range(len(concrete.shape) - 1, -1, -1):
+            strides.append(concrete.strides[i])
+            backstrides.append(concrete.backstrides[i])
+            shape.append(concrete.shape[i])
+        return space.wrap(NDimSlice(concrete, new_sig, self.start, strides[:],
                            backstrides[:], shape[:]))
 
     def descr_get_flatiter(self, space):
@@ -821,7 +830,7 @@
         return self.dtype
 
     def getitem(self, item):
-        return self.value
+        raise NotImplementedError
 
     def eval(self, iter):
         return self.value
@@ -832,6 +841,7 @@
     def to_str(self, space, comma, builder, indent=' ', use_ellipsis=False):
         builder.append(self.dtype.str_format(self.value))
 
+
 class VirtualArray(BaseArray):
     """
     Class for representing virtual arrays, such as binary ops or ufuncs
@@ -1198,7 +1208,7 @@
 )
 
 def descr_new_flatiter(space, w_object):
-    assert isinstance(w_object,BaseArray)
+    assert isinstance(w_object, BaseArray)
     i = FlatIterator(w_object)
     return i
 
@@ -1215,11 +1225,11 @@
 
     def descr_next(self, space):
         if self.iter.done():
-            raise OperationError(space.w_StopIteration,space.wrap(''))
+            raise OperationError(space.w_StopIteration, space.wrap(''))
         retVal = self.arr.eval(self.iter)
         self.iter = self.iter.next(self.shapelen)
         return retVal.wrap(space)
-            
+
 
 FlatIterator.typedef = TypeDef(
     'flatiter',
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
@@ -214,9 +214,11 @@
         from numpypy import array
         a = array(range(5))
         assert a[3] == 3
-        a = array(1)
-        assert a[0] == 1
-        assert a.shape == ()
+        #Wrong on two counts: numpy does not allow assigning to Scalar,
+        # and Scalar.shape is not a test of iterator_init, is it?
+        #a = array(1)
+        #assert a[0] == 1
+        #assert a.shape == ()
 
     def test_getitem(self):
         from numpypy import array
@@ -304,7 +306,10 @@
     def test_scalar(self):
         from numpypy import array
         a = array(3)
-        assert a[0] == 3
+        #assert a[0] == 3
+        raises(IndexError, "a[0]")
+        assert a.size == 1
+        assert a.shape == ()
 
     def test_len(self):
         from numpypy import array