[pypy-commit] pypy numpypy-ellipse-indexing: implement, tests pass

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: numpypy-ellipse-indexing
Changeset: r64061:6d277aa40577
Date: 2013-05-14 09:15 +0300
http://bitbucket.org/pypy/pypy/changeset/6d277aa40577/

Log:	implement, tests pass

diff --git a/pypy/module/micronumpy/arrayimpl/concrete.py b/pypy/module/micronumpy/arrayimpl/concrete.py
--- a/pypy/module/micronumpy/arrayimpl/concrete.py
+++ b/pypy/module/micronumpy/arrayimpl/concrete.py
@@ -205,7 +205,7 @@
         if (space.isinstance_w(w_idx, space.w_int) or
             space.isinstance_w(w_idx, space.w_slice)):
             return Chunks([Chunk(*space.decode_index4(w_idx, self.get_shape()[0]))])
-        elif space.is_w(w_idx, space.w_None):
+        elif space.is_w(w_idx, space.w_None) or isinstance(w_idx, Ellipsis):
             return Chunks([NewAxisChunk()])
         result = []
         i = 0
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
@@ -20,6 +20,7 @@
 from rpython.rlib import jit
 from rpython.rlib.rstring import StringBuilder
 from pypy.module.micronumpy.arrayimpl.base import BaseArrayImplementation
+from pypy.interpreter.special import Ellipsis
 
 def _find_shape(space, w_size):
     if space.is_none(w_size):
@@ -167,6 +168,8 @@
                                prefix)
 
     def descr_getitem(self, space, w_idx):
+        if isinstance(w_idx, Ellipsis):
+            return self
         if (isinstance(w_idx, W_NDimArray) and
             w_idx.get_dtype().is_bool_type()):
             return self.getitem_filter(space, w_idx)
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
@@ -1661,11 +1661,17 @@
         assert (b == [20, 1, 21, 3, 4]).all()
         raises(ValueError, "array([1, 2])[array([True, False, True])] = [1, 2, 3]")
 
-    def test_ellipse_index_setitem(self):
+    def test_ellipse_index(self):
         from numpypy import arange
-        b = arange(5)
+        b = arange(24).reshape(2,3,4)
         b[...] = 100
         assert (b == 100).all()
+        assert b.shape == (2, 3, 4)
+        b[...] = [10, 20, 30, 40]
+        assert (b[:,:,0] == 10).all()
+        assert (b[0,0,:] == [10, 20, 30, 40]).all()
+        assert b.shape == b[...].shape
+        assert (b == b[...]).all()
 
     def test_weakref(self):
         import _weakref