[pypy-commit] pypy numpy-multidim: tentative checkin

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-multidim
Changeset: r48492:4dd7d695536a
Date: 2011-10-27 00:04 +0200
http://bitbucket.org/pypy/pypy/changeset/4dd7d695536a/

Log:	tentative checkin

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
@@ -223,17 +223,32 @@
         concrete = self.get_concrete()
         return space.wrap("[" + " ".join(concrete._getnums(True)) + "]")
 
+    def item_at_index(self, index, space):
+        # we assume C ordering for now
+        item = 0
+        for i in range(len(index)):
+            if i != 0:
+                item *= self.shape[i]
+            if index[i] >= self.shape[i]:
+                raise OperationError(space.w_IndexError,
+                                     space.wrap("index (%d) out of range (0<=index<%d" % (index[i], self.shape[i])))
+            item += index[i]
+        return item
+
     def descr_getitem(self, space, w_idx):
         # TODO: indexing by arrays and lists
         if space.isinstance_w(w_idx, space.w_tuple):
+            # or any other sequence actually
             length = space.len_w(w_idx)
             if length == 0:
                 return space.wrap(self)
-            if length > 1: # only one dimension for now.
+            if length > len(self.shape):
                 raise OperationError(space.w_IndexError,
                                      space.wrap("invalid index"))
-            w_idx = space.getitem(w_idx, space.wrap(0))
-        start, stop, step, slice_length = space.decode_index4(w_idx, self.find_size())
+            indices = [space.int_w(w_item) for w_item in space.fixedview(w_idx)]
+            item = self.item_at_index(indices, space)
+            return self.get_concrete().eval(item).wrap(space)
+        start, stop, step, slice_length = space.decode_index4(w_idx, self.shape[0])
         if step == 0:
             # Single index
             return self.get_concrete().eval(start).wrap(space)
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
@@ -611,7 +611,17 @@
         assert numpy.zeros((2, 2)).shape == (2,2)
         assert numpy.zeros((3, 1, 2)).shape == (3, 1, 2)
         assert len(numpy.zeros((3, 1, 2))) == 3
-        
+
+    def test_getsetitem(self):
+        import numpy
+        a = numpy.zeros((2, 3, 1))
+        raises(IndexError, a.__getitem__, (0, 0, 0, 0))
+        raises(IndexError, a.__getitem__, (3,))
+        raises(IndexError, a.__getitem__, (1, 3))
+        assert a[1, 1, 0] == 0
+        a[1, 2, 0] = 3
+        assert a[1, 2, 0] == 3
+        assert a[1, 1, 0] == 0
 
 class AppTestSupport(object):
     def setup_class(cls):