[pypy-commit] pypy default: merge

fijal noreply at buildbot.pypy.org
Mon Jul 20 18:31:19 CEST 2015


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r78613:4b6cfd24db0b
Date: 2015-07-20 18:31 +0200
http://bitbucket.org/pypy/pypy/changeset/4b6cfd24db0b/

Log:	merge

diff --git a/pypy/doc/whatsnew-head.rst b/pypy/doc/whatsnew-head.rst
--- a/pypy/doc/whatsnew-head.rst
+++ b/pypy/doc/whatsnew-head.rst
@@ -43,3 +43,6 @@
 Improve compatibility with numpy dtypes; handle offsets to create unions,
 fix str() and repr(), allow specifying itemsize, metadata and titles, add flags,
 allow subclassing dtype
+
+.. branch: indexing
+Refactor array indexing to support ellipses.
diff --git a/pypy/module/micronumpy/arrayops.py b/pypy/module/micronumpy/arrayops.py
--- a/pypy/module/micronumpy/arrayops.py
+++ b/pypy/module/micronumpy/arrayops.py
@@ -5,7 +5,7 @@
 from pypy.module.micronumpy.base import convert_to_array, W_NDimArray
 from pypy.module.micronumpy.converters import clipmode_converter
 from pypy.module.micronumpy.strides import (
-    Chunk, Chunks, shape_agreement, shape_agreement_multiple)
+    Chunk, new_view, shape_agreement, shape_agreement_multiple)
 from .casting import find_binop_result_dtype, find_result_type
 
 
@@ -148,7 +148,8 @@
             continue
         chunks[axis] = Chunk(axis_start, axis_start + arr.get_shape()[axis], 1,
                              arr.get_shape()[axis])
-        Chunks(chunks).apply(space, res).implementation.setslice(space, arr)
+        view = new_view(space, res, chunks)
+        view.implementation.setslice(space, arr)
         axis_start += arr.get_shape()[axis]
     return res
 
@@ -162,8 +163,9 @@
         shape = [arr.get_shape()[0] * repeats]
         w_res = W_NDimArray.from_shape(space, shape, arr.get_dtype(), w_instance=arr)
         for i in range(repeats):
-            Chunks([Chunk(i, shape[0] - repeats + i, repeats,
-                 orig_size)]).apply(space, w_res).implementation.setslice(space, arr)
+            chunks = [Chunk(i, shape[0] - repeats + i, repeats, orig_size)]
+            view = new_view(space, w_res, chunks)
+            view.implementation.setslice(space, arr)
     else:
         axis = space.int_w(w_axis)
         shape = arr.get_shape()[:]
@@ -174,7 +176,8 @@
         for i in range(repeats):
             chunks[axis] = Chunk(i, shape[axis] - repeats + i, repeats,
                                  orig_size)
-            Chunks(chunks).apply(space, w_res).implementation.setslice(space, arr)
+            view = new_view(space, w_res, chunks)
+            view.implementation.setslice(space, arr)
     return w_res
 
 
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
@@ -1,7 +1,7 @@
 from pypy.interpreter.error import OperationError, oefmt
 from rpython.rlib import jit, rgc
 from rpython.rlib.buffer import Buffer
-from rpython.rlib.debug import make_sure_not_resized, debug_print
+from rpython.rlib.debug import make_sure_not_resized
 from rpython.rlib.rawstorage import alloc_raw_storage, free_raw_storage, \
     raw_storage_getitem, raw_storage_setitem, RAW_STORAGE
 from rpython.rtyper.lltypesystem import rffi, lltype, llmemory
@@ -9,13 +9,12 @@
 from pypy.module.micronumpy.base import convert_to_array, W_NDimArray, \
     ArrayArgumentException, W_NumpyObject
 from pypy.module.micronumpy.iterators import ArrayIter
-from pypy.module.micronumpy.strides import (Chunk, Chunks, NewAxisChunk,
-    RecordChunk, calc_strides, calc_new_strides, shape_agreement,
+from pypy.module.micronumpy.strides import (
+    IntegerChunk, SliceChunk, NewAxisChunk, EllipsisChunk, new_view,
+    calc_strides, calc_new_strides, shape_agreement,
     calculate_broadcast_strides, calc_backstrides, calc_start, is_c_contiguous,
     is_f_contiguous)
 from rpython.rlib.objectmodel import keepalive_until_here
-from rpython.rtyper.annlowlevel import cast_gcref_to_instance
-from pypy.interpreter.baseobjspace import W_Root
 
 class BaseConcreteArray(object):
     _immutable_fields_ = ['dtype?', 'storage', 'start', 'size', 'shape[*]',
@@ -204,6 +203,8 @@
                 if (isinstance(w_item, W_NDimArray) or
                     space.isinstance_w(w_item, space.w_list)):
                     raise ArrayArgumentException
+                elif space.is_w(w_item, space.w_Ellipsis):
+                    raise IndexError
             return self._lookup_by_index(space, view_w)
         if shape_len == 0:
             raise oefmt(space.w_IndexError, "too many indices for array")
@@ -215,39 +216,47 @@
     @jit.unroll_safe
     def _prepare_slice_args(self, space, w_idx):
         if space.isinstance_w(w_idx, space.w_str):
-            idx = space.str_w(w_idx)
-            dtype = self.dtype
-            if not dtype.is_record():
-                raise oefmt(space.w_IndexError, "only integers, slices (`:`), "
-                    "ellipsis (`...`), numpy.newaxis (`None`) and integer or "
-                    "boolean arrays are valid indices")
-            elif idx not in dtype.fields:
-                raise oefmt(space.w_ValueError, "field named %s not found", idx)
-            return RecordChunk(idx)
-        elif (space.isinstance_w(w_idx, space.w_int) or
-                space.isinstance_w(w_idx, space.w_slice)):
+            raise oefmt(space.w_IndexError, "only integers, slices (`:`), "
+                "ellipsis (`...`), numpy.newaxis (`None`) and integer or "
+                "boolean arrays are valid indices")
+        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 Chunks([Chunk(*space.decode_index4(w_idx, self.get_shape()[0]))])
+            return [SliceChunk(w_idx), EllipsisChunk()]
+        elif space.isinstance_w(w_idx, space.w_int):
+            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 Chunks([Chunk(*space.decode_index4(w_idx, self.get_shape()[0]))])
+            return [IntegerChunk(w_idx), EllipsisChunk()]
         elif space.is_w(w_idx, space.w_None):
-            return Chunks([NewAxisChunk()])
+            return [NewAxisChunk(), EllipsisChunk()]
         result = []
         i = 0
+        has_ellipsis = False
         for w_item in space.fixedview(w_idx):
-            if space.is_w(w_item, space.w_None):
+            if space.is_w(w_item, space.w_Ellipsis):
+                if has_ellipsis:
+                    # in CNumPy, this is only a deprecation warning
+                    raise oefmt(space.w_ValueError,
+                        "an index can only have a single Ellipsis (`...`); "
+                        "replace all but one with slices (`:`).")
+                result.append(EllipsisChunk())
+                has_ellipsis = True
+            elif space.is_w(w_item, space.w_None):
                 result.append(NewAxisChunk())
+            elif space.isinstance_w(w_item, space.w_slice):
+                result.append(SliceChunk(w_item))
+                i += 1
             else:
-                result.append(Chunk(*space.decode_index4(w_item,
-                                                         self.get_shape()[i])))
+                result.append(IntegerChunk(w_item))
                 i += 1
-        return Chunks(result)
+        if not has_ellipsis:
+            result.append(EllipsisChunk())
+        return result
 
     def descr_getitem(self, space, orig_arr, w_index):
         try:
@@ -256,7 +265,7 @@
         except IndexError:
             # not a single result
             chunks = self._prepare_slice_args(space, w_index)
-            return chunks.apply(space, orig_arr)
+            return new_view(space, orig_arr, chunks)
 
     def descr_setitem(self, space, orig_arr, w_index, w_value):
         try:
@@ -265,7 +274,7 @@
         except IndexError:
             w_value = convert_to_array(space, w_value)
             chunks = self._prepare_slice_args(space, w_index)
-            view = chunks.apply(space, orig_arr)
+            view = new_view(space, orig_arr, chunks)
             view.implementation.setslice(space, w_value)
 
     def transpose(self, orig_array, axes=None):
diff --git a/pypy/module/micronumpy/ndarray.py b/pypy/module/micronumpy/ndarray.py
--- a/pypy/module/micronumpy/ndarray.py
+++ b/pypy/module/micronumpy/ndarray.py
@@ -18,8 +18,9 @@
 from pypy.module.micronumpy.converters import multi_axis_converter, \
     order_converter, shape_converter, searchside_converter
 from pypy.module.micronumpy.flagsobj import W_FlagsObject
-from pypy.module.micronumpy.strides import get_shape_from_iterable, \
-    shape_agreement, shape_agreement_multiple, is_c_contiguous, is_f_contiguous
+from pypy.module.micronumpy.strides import (
+    get_shape_from_iterable, shape_agreement, shape_agreement_multiple,
+    is_c_contiguous, is_f_contiguous, calc_strides, new_view)
 from pypy.module.micronumpy.casting import can_cast_array
 
 
@@ -178,7 +179,7 @@
         if iter_shape is None:
             # w_index is a list of slices, return a view
             chunks = self.implementation._prepare_slice_args(space, w_index)
-            return chunks.apply(space, self)
+            return new_view(space, self, chunks)
         shape = res_shape + self.get_shape()[len(indexes):]
         w_res = W_NDimArray.from_shape(space, shape, self.get_dtype(),
                                        self.get_order(), w_instance=self)
@@ -194,7 +195,7 @@
         if iter_shape is None:
             # w_index is a list of slices
             chunks = self.implementation._prepare_slice_args(space, w_index)
-            view = chunks.apply(space, self)
+            view = new_view(space, self, chunks)
             view.implementation.setslice(space, val_arr)
             return
         if support.product(iter_shape) == 0:
@@ -203,6 +204,10 @@
                                prefix)
 
     def descr_getitem(self, space, w_idx):
+        if self.get_dtype().is_record():
+            if space.isinstance_w(w_idx, space.w_str):
+                idx = space.str_w(w_idx)
+                return self.getfield(space, idx)
         if space.is_w(w_idx, space.w_Ellipsis):
             return self
         elif isinstance(w_idx, W_NDimArray) and w_idx.get_dtype().is_bool():
@@ -229,6 +234,13 @@
         self.implementation.setitem_index(space, index_list, w_value)
 
     def descr_setitem(self, space, w_idx, w_value):
+        if self.get_dtype().is_record():
+            if space.isinstance_w(w_idx, space.w_str):
+                idx = space.str_w(w_idx)
+                view = self.getfield(space, idx)
+                w_value = convert_to_array(space, w_value)
+                view.implementation.setslice(space, w_value)
+                return
         if space.is_w(w_idx, space.w_Ellipsis):
             self.implementation.setslice(space, convert_to_array(space, w_value))
             return
@@ -241,6 +253,28 @@
         except ArrayArgumentException:
             self.setitem_array_int(space, w_idx, w_value)
 
+    def getfield(self, space, field):
+        dtype = self.get_dtype()
+        if field not in dtype.fields:
+            raise oefmt(space.w_ValueError, "field named %s not found", field)
+        arr = self.implementation
+        ofs, subdtype = arr.dtype.fields[field][:2]
+        # ofs only changes start
+        # create a view of the original array by extending
+        # the shape, strides, backstrides of the array
+        strides, backstrides = calc_strides(subdtype.shape,
+                                            subdtype.subdtype, arr.order)
+        final_shape = arr.shape + subdtype.shape
+        final_strides = arr.get_strides() + strides
+        final_backstrides = arr.get_backstrides() + backstrides
+        final_dtype = subdtype
+        if subdtype.subdtype:
+            final_dtype = subdtype.subdtype
+        return W_NDimArray.new_slice(space, arr.start + ofs, final_strides,
+                                     final_backstrides,
+                                     final_shape, arr, self, final_dtype)
+
+
     def descr_delitem(self, space, w_idx):
         raise OperationError(space.w_ValueError, space.wrap(
             "cannot delete array elements"))
@@ -1298,7 +1332,6 @@
 def descr_new_array(space, w_subtype, w_shape, w_dtype=None, w_buffer=None,
                     offset=0, w_strides=None, w_order=None):
     from pypy.module.micronumpy.concrete import ConcreteArray
-    from pypy.module.micronumpy.strides import calc_strides
     dtype = space.interp_w(descriptor.W_Dtype, space.call_function(
         space.gettypefor(descriptor.W_Dtype), w_dtype))
     shape = shape_converter(space, w_shape, dtype)
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
@@ -10,78 +10,92 @@
     pass
 
 
-class RecordChunk(BaseChunk):
-    def __init__(self, name):
-        self.name = name
-
-    def apply(self, space, orig_arr):
-        arr = orig_arr.implementation
-        ofs, subdtype = arr.dtype.fields[self.name][:2]
-        # ofs only changes start
-        # create a view of the original array by extending
-        # the shape, strides, backstrides of the array
-        strides, backstrides = calc_strides(subdtype.shape,
-                                            subdtype.subdtype, arr.order)
-        final_shape = arr.shape + subdtype.shape
-        final_strides = arr.get_strides() + strides
-        final_backstrides = arr.get_backstrides() + backstrides
-        final_dtype = subdtype
-        if subdtype.subdtype:
-            final_dtype = subdtype.subdtype
-        return W_NDimArray.new_slice(space, arr.start + ofs, final_strides,
-                                     final_backstrides,
-                                     final_shape, arr, orig_arr, final_dtype)
-
-
-class Chunks(BaseChunk):
-    def __init__(self, l):
-        self.l = l
-
-    @jit.unroll_safe
-    def extend_shape(self, old_shape):
-        shape = []
-        i = -1
-        for i, c in enumerate_chunks(self.l):
-            if c.step != 0:
-                shape.append(c.lgt)
-        s = i + 1
-        assert s >= 0
-        return shape[:] + old_shape[s:]
-
-    def apply(self, space, orig_arr):
-        arr = orig_arr.implementation
-        shape = self.extend_shape(arr.shape)
-        r = calculate_slice_strides(arr.shape, arr.start, arr.get_strides(),
-                                    arr.get_backstrides(), self.l)
-        _, start, strides, backstrides = r
-        return W_NDimArray.new_slice(space, start, strides[:], backstrides[:],
-                                     shape[:], arr, orig_arr)
-
-
 class Chunk(BaseChunk):
-    axis_step = 1
+    input_dim = 1
 
     def __init__(self, start, stop, step, lgt):
         self.start = start
         self.stop = stop
         self.step = step
         self.lgt = lgt
+        if self.step == 0:
+            self.out_dim = 0
+        else:
+            self.out_dim = 1
+
+    def compute(self, space, base_length, base_stride):
+        stride = base_stride * self.step
+        backstride = base_stride * max(0, self.lgt - 1) * self.step
+        return self.start, self.lgt, stride, backstride
 
     def __repr__(self):
         return 'Chunk(%d, %d, %d, %d)' % (self.start, self.stop, self.step,
                                           self.lgt)
 
+class IntegerChunk(BaseChunk):
+    input_dim = 1
+    out_dim = 0
+    def __init__(self, w_idx):
+        self.w_idx = w_idx
+
+    def compute(self, space, base_length, base_stride):
+        start, _, _, _ = space.decode_index4(self.w_idx, base_length)
+        return start, 0, 0, 0
+
+
+class SliceChunk(BaseChunk):
+    input_dim = 1
+    out_dim = 1
+
+    def __init__(self, w_slice):
+        self.w_slice = w_slice
+
+    def compute(self, space, base_length, base_stride):
+        start, stop, step, length = space.decode_index4(self.w_slice, base_length)
+        stride = base_stride * step
+        backstride = base_stride * max(0, length - 1) * step
+        return start, length, stride, backstride
 
 class NewAxisChunk(Chunk):
-    start = 0
-    stop = 1
-    step = 1
-    lgt = 1
-    axis_step = 0 # both skip this axis in calculate_slice_strides and set stride => 0
+    input_dim = 0
+    out_dim = 1
 
     def __init__(self):
         pass
 
+    def compute(self, space, base_length, base_stride):
+        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
+    r = calculate_slice_strides(space, arr.shape, arr.start, arr.get_strides(),
+                                arr.get_backstrides(), chunks)
+    shape, start, strides, backstrides = r
+    return W_NDimArray.new_slice(space, start, strides[:], backstrides[:],
+                                 shape[:], arr, w_arr)
+
+ at jit.unroll_safe
+def _extend_shape(old_shape, chunks):
+    shape = []
+    i = -1
+    for i, c in enumerate_chunks(chunks):
+        if c.out_dim > 0:
+            shape.append(c.lgt)
+    s = i + 1
+    assert s >= 0
+    return shape[:] + old_shape[s:]
+
 
 class BaseTransform(object):
     pass
@@ -103,41 +117,56 @@
     result = []
     i = -1
     for chunk in chunks:
-        i += chunk.axis_step
+        i += chunk.input_dim
         result.append((i, chunk))
     return result
 
 
- at jit.look_inside_iff(lambda shape, start, strides, backstrides, chunks:
+ at jit.look_inside_iff(lambda space, shape, start, strides, backstrides, chunks:
                      jit.isconstant(len(chunks)))
-def calculate_slice_strides(shape, start, strides, backstrides, 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:
-        if chunk.step != 0:
-            size += 1
-    rstrides = [0] * size
-    rbackstrides = [0] * size
+        used_dims += chunk.input_dim
+        size += chunk.out_dim
+    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
-    j = 0
-    for i, chunk in enumerate_chunks(chunks):
-        try:
-            s_i = strides[i]
-        except IndexError:
+    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:
+        if isinstance(chunk, NewAxisChunk):
+            rshape[j] = 1
+            j += 1
             continue
-        if chunk.step != 0:
-            rstrides[j] = s_i * chunk.step * chunk.axis_step
-            rbackstrides[j] = s_i * max(0, chunk.lgt - 1) * chunk.step
-            rshape[j] = chunk.lgt
-            j += 1
-        rstart += s_i * chunk.start
-    # add a reminder
-    s = i + 1
-    assert s >= 0
-    rstrides += strides[s:]
-    rbackstrides += backstrides[s:]
-    rshape += shape[s:]
+        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
+            j += chunk.out_dim
+        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, Chunks
+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,7 +22,9 @@
 
 
 def create_slice(space, a, chunks):
-    return Chunks(chunks).apply(space, W_NDimArray(a)).implementation
+    if not any(isinstance(c, EllipsisChunk) for c in chunks):
+        chunks.append(EllipsisChunk())
+    return new_view(space, W_NDimArray(a), chunks).implementation
 
 
 def create_array(*args, **kwargs):
@@ -2488,6 +2490,13 @@
         assert b.shape == b[...].shape
         assert (b == b[...]).all()
 
+        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
         r = np.ones(3)


More information about the pypy-commit mailing list