[pypy-commit] pypy indexing: implement ellipsis indexing

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: indexing
Changeset: r78588:5663f55d1d87
Date: 2015-07-17 20:02 +0100
http://bitbucket.org/pypy/pypy/changeset/5663f55d1d87/

Log:	implement ellipsis indexing

diff --git a/pypy/module/micronumpy/concrete.py b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -222,18 +222,18 @@
         if space.isinstance_w(w_idx, space.w_slice):
             if len(self.get_shape()) == 0:
                 raise oefmt(space.w_ValueError, "cannot slice a 0-d array")
-            return [SliceChunk(w_idx)]
+            return [SliceChunk(w_idx), EllipsisChunk()]
         elif space.isinstance_w(w_idx, space.w_int):
-            return [IntegerChunk(w_idx)]
+            return [IntegerChunk(w_idx), EllipsisChunk()]
         elif isinstance(w_idx, W_NDimArray) and w_idx.is_scalar():
             w_idx = w_idx.get_scalar_value().item(space)
             if not space.isinstance_w(w_idx, space.w_int) and \
                     not space.isinstance_w(w_idx, space.w_bool):
                 raise OperationError(space.w_IndexError, space.wrap(
                     "arrays used as indices must be of integer (or boolean) type"))
-            return [IntegerChunk(w_idx)]
+            return [IntegerChunk(w_idx), EllipsisChunk()]
         elif space.is_w(w_idx, space.w_None):
-            return [NewAxisChunk()]
+            return [NewAxisChunk(), EllipsisChunk()]
         result = []
         i = 0
         has_ellipsis = False
@@ -254,6 +254,8 @@
             else:
                 result.append(IntegerChunk(w_item))
                 i += 1
+        if not has_ellipsis:
+            result.append(EllipsisChunk())
         return result
 
     def descr_getitem(self, space, orig_arr, w_index):
diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -59,12 +59,7 @@
         backstride = base_stride * max(0, length - 1) * step
         return start, length, stride, backstride
 
-
 class NewAxisChunk(Chunk):
-    start = 0
-    stop = 1
-    step = 1
-    lgt = 1
     input_dim = 0
     out_dim = 1
 
@@ -75,9 +70,15 @@
         return 0, 1, 0, 0
 
 class EllipsisChunk(BaseChunk):
+    input_dim = 0
+    out_dim = 0
     def __init__(self):
         pass
 
+    def compute(self, space, base_length, base_stride):
+        backstride = base_stride * max(0, base_length - 1)
+        return 0, base_length, base_stride, backstride
+
 
 def new_view(space, w_arr, chunks):
     arr = w_arr.implementation
@@ -127,38 +128,48 @@
 @jit.look_inside_iff(lambda space, shape, start, strides, backstrides, chunks:
                      jit.isconstant(len(chunks)))
 def calculate_slice_strides(space, shape, start, strides, backstrides, chunks):
+    """
+    Note: `chunks` must contain exactly one EllipsisChunk object.
+    """
     size = 0
+    used_dims = 0
     for chunk in chunks:
+        used_dims += chunk.input_dim
         size += chunk.out_dim
-    rstrides = [0] * size
-    rbackstrides = [0] * size
+    if used_dims > len(shape):
+        raise oefmt(space.w_IndexError, "too many indices for array")
+    else:
+        extra_dims = len(shape) - used_dims
+    rstrides = [0] * (size + extra_dims)
+    rbackstrides = [0] * (size + extra_dims)
     rstart = start
-    rshape = [0] * size
-    i = -1  # index of the current dimension in the input array
+    rshape = [0] * (size + extra_dims)
+    rstart = start
+    i = 0  # index of the current dimension in the input array
     j = 0  # index of the current dimension in the result view
     for chunk in chunks:
-        i += chunk.input_dim
         if isinstance(chunk, NewAxisChunk):
             rshape[j] = 1
             j += 1
             continue
-        try:
-            s_i = strides[i]
-        except IndexError:
+        elif isinstance(chunk, EllipsisChunk):
+            for k in range(extra_dims):
+                start, length, stride, backstride = chunk.compute(
+                        space, shape[i], strides[i])
+                rshape[j] = length
+                rstrides[j] = stride
+                rbackstrides[j] = backstride
+                j += 1
+                i += 1
             continue
         start, length, stride, backstride = chunk.compute(space, shape[i], strides[i])
         if chunk.out_dim == 1:
+            rshape[j] = length
             rstrides[j] = stride
             rbackstrides[j] = backstride
-            rshape[j] = length
             j += chunk.out_dim
-        rstart += s_i * start
-    # add a reminder
-    s = i + 1
-    assert s >= 0
-    rstrides += strides[s:]
-    rbackstrides += backstrides[s:]
-    rshape += shape[s:]
+        rstart += strides[i] * start
+        i += chunk.input_dim
     return rshape, rstart, rstrides, rbackstrides
 
 
diff --git a/pypy/module/micronumpy/test/test_ndarray.py b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -4,7 +4,7 @@
 
 from pypy.conftest import option
 from pypy.module.micronumpy.appbridge import get_appbridge_cache
-from pypy.module.micronumpy.strides import Chunk, new_view
+from pypy.module.micronumpy.strides import Chunk, new_view, EllipsisChunk
 from pypy.module.micronumpy.ndarray import W_NDimArray
 from pypy.module.micronumpy.test.test_base import BaseNumpyAppTest
 
@@ -22,6 +22,8 @@
 
 
 def create_slice(space, a, chunks):
+    if not any(isinstance(c, EllipsisChunk) for c in chunks):
+        chunks.append(EllipsisChunk())
     return new_view(space, W_NDimArray(a), chunks).implementation
 
 
@@ -2490,6 +2492,9 @@
         a = np.arange(6).reshape(2, 3)
         if '__pypy__' in sys.builtin_module_names:
             raises(ValueError, "a[..., ...]")
+        b = a [..., 0]
+        assert (b == [0, 3]).all()
+        assert b.base is a
 
     def test_empty_indexing(self):
         import numpy as np