ENH: Expose prepare_index_tuple as index_tricks.as_index_tuple

Approaches #8275
numpy · eric-wieser · Nov 14, 2016 · Dec 5, 2016 · Dec 21, 2016 · Dec 21, 2016
commit a86cc8c63f417431c27b3c545b784aad9824319a
diff --git a/doc/release/1.13.0-notes.rst b/doc/release/1.13.0-notes.rst
@@ -159,6 +159,11 @@ being iterated over.
 For consistency with ``ndarray`` and ``broadcast``, ``d.ndim`` is a shorthand
 for ``len(d.shape)``.
 
+``as_index_tuple`` function in ``index_tricks``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Exposes the internal normalization that happens when indexing with non-tuple
+objects to convert them into tuples. Useful for correctly overriding or
+wrapping ndarray.__getitem__
 
 Improvements
 ============

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
@@ -5403,6 +5403,22 @@ def luf(lamdaexpr, *args, **kwargs):
 
     """)
 
+add_newdoc('numpy.core.multiarray', 'as_index_tuple',
+    """
+    as_index_tuple(index)
+
+    Normalizes an index argument, like that passed to `__getitem__`, into a tuple.
+    Follows the invariant that `x[index]` is identical to `x[as_index_tuple(index)].
+
+    Examples
+    --------
+    >>> as_index_tuple(1)
+    (1,)
+    >>> as_index_tuple([1, 2, None])
+    (1, 2, None)
+    >>> as_index_tuple([1, 2, 3])
+    ([1, 2, 3])
+    """)
 
 ##############################################################################
 #

diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
@@ -284,6 +284,40 @@ prepare_index_tuple(PyObject *index, PyObject **result)
     }
 }
 
+/**
+ * Expose prepare_index_tuple to python code
+ */
+NPY_NO_EXPORT PyObject *
+as_index_tuple(PyObject *NPY_UNUSED(self), PyObject *args)
+{
+    PyObject *obj;
+    PyObject *result;
+    PyObject *prepared[NPY_MAXDIMS*2];
+    npy_intp i, n;
+
+    if (!PyArg_ParseTuple(args, "O", &obj)) {
+        return NULL;
+    }
+    n = prepare_index_tuple(obj, prepared);
+    if (n < 0) {
+        return NULL;
+    }
+
+    result = PyTuple_New(n);
+    if (result == NULL) {
+        return NULL;
+    }
+
+    for (i = 0; i < n; i++) {
+        PyObject *val = prepared[i];
+        Py_INCREF(val);
+        PyTuple_SET_ITEM(result, i, val);
+    }
+
+    return result;
+}
+
+
 /**
  * Prepare an npy_index_object from the python slicing object.
  *

diff --git a/numpy/core/src/multiarray/mapping.h b/numpy/core/src/multiarray/mapping.h
@@ -47,6 +47,9 @@ array_subscript(PyArrayObject *self, PyObject *op);
 NPY_NO_EXPORT int
 array_assign_item(PyArrayObject *self, Py_ssize_t i, PyObject *v);
 
+NPY_NO_EXPORT PyObject *
+as_index_tuple(PyObject *NPY_UNUSED(self), PyObject *args);
+
 /*
  * Prototypes for Mapping calls --- not part of the C-API
  * because only useful as part of a getitem call.

diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
@@ -60,6 +60,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 #include "templ_common.h" /* for npy_mul_with_overflow_intp */
 #include "compiled_base.h"
 #include "mem_overlap.h"
+#include "mapping.h" /* for as_index_tuple */
 
 /* Only here for API compatibility */
 NPY_NO_EXPORT PyTypeObject PyBigArray_Type;
@@ -4293,6 +4294,8 @@ static struct PyMethodDef array_module_methods[] = {
         METH_VARARGS | METH_KEYWORDS, NULL},
     {"normalize_axis_index", (PyCFunction)normalize_axis_index,
         METH_VARARGS | METH_KEYWORDS, NULL},
+    {"as_index_tuple", (PyCFunction)as_index_tuple,
+        METH_VARARGS, NULL},
     {NULL, NULL, 0, NULL}                /* sentinel */
 };
 

diff --git a/numpy/core/tests/test_indexing.py b/numpy/core/tests/test_indexing.py
@@ -511,6 +511,98 @@ def test_indexing_array_negative_strides(self):
         arr[slices] = 10
         assert_array_equal(arr, 10.)
 
+    def test_as_index_tuple(self):
+        from numpy.core.multiarray import as_index_tuple
+
+        arr = np.array([1])
+        sl = np.s_[:]
+        ell = Ellipsis
+        obj = object()
+
+        # scalars are wrapped in a 1-tuple
+        assert_equal(as_index_tuple(1),    (1,))
+        assert_equal(as_index_tuple(ell),  (ell,))
+        assert_equal(as_index_tuple(None), (None,))
+        assert_equal(as_index_tuple(sl),   (sl,))
+        assert_equal(as_index_tuple(arr),  (arr,))
+        assert_equal(as_index_tuple(obj),  (obj,))
+
+        # tuples are untouched
+        assert_equal(as_index_tuple((1, 2, 3)),    (1, 2, 3))
+        assert_equal(as_index_tuple((1, 2, ell)),  (1, 2, ell))
+        assert_equal(as_index_tuple((1, 2, None)), (1, 2, None))
+        assert_equal(as_index_tuple((1, 2, sl)),   (1, 2, sl))
+        assert_equal(as_index_tuple((1, 2, arr)),  (1, 2, arr))
+
+        # sequences of scalars are wrapped
+        assert_equal(as_index_tuple([1, 2, 3]), ([1, 2, 3],))
+
+        # sequences containing slice objects or ellipses are tuple-ified
+        assert_equal(as_index_tuple([1, 2, ell]),  (1, 2, ell))
+        assert_equal(as_index_tuple([1, 2, None]), (1, 2, None))
+        assert_equal(as_index_tuple([1, 2, sl]),   (1, 2, sl))
+        assert_equal(as_index_tuple([1, 2, arr]),  (1, 2, arr))
+
+        # unless they are >= np.MAXDIMS, in which case they are always wrapped
+        nd = np.MAXDIMS
+        assert_equal(as_index_tuple(nd * [1]),    (nd * [1],))
+        assert_equal(as_index_tuple(nd * [ell]),  (nd * [ell],))
+        assert_equal(as_index_tuple(nd * [None]), (nd * [None],))
+        assert_equal(as_index_tuple(nd * [sl]),   (nd * [sl],))
+        assert_equal(as_index_tuple(nd * [arr]),  (nd * [arr],))
+
+    def test_as_index_tuple_broken_getitem(self):
+        from numpy.core.multiarray import as_index_tuple
+
+        # test sequences with a broken __getitem__
+        def make_broken_sequence(base, items):
+            class Broken(base):
+                def __len__(self):
+                    return len(items)
+                def __getitem__(self, i):
+                    val = items[i]
+                    if isinstance(val, Exception):
+                        raise val
+                    return val
+            return Broken()
+
+        # error comes first, so just treat as a scalar
+        idx = make_broken_sequence(object, [1, ValueError(), None])
+        assert_raises(ValueError, operator.getitem, idx, 1)
+        assert_equal(as_index_tuple(idx),  (idx,))
+
+        # none comes first, so commit to making the tuple
+        idx = make_broken_sequence(object, [1, None, ValueError()])
+        assert_raises(ValueError, operator.getitem, idx, 2)
+        assert_raises(ValueError, as_index_tuple, idx)
+
+        # tuples subclasses error in both cases
+        idx = make_broken_sequence(tuple, [1, ValueError(), None])
+        assert_raises(ValueError, operator.getitem, idx, 1)
+        assert_raises(ValueError, as_index_tuple, idx)
+
+        idx = make_broken_sequence(tuple, [1, None, ValueError()])
+        assert_raises(ValueError, operator.getitem, idx, 2)
+        assert_raises(ValueError, as_index_tuple, idx)
+
+    def test_as_index_tuple_broken_len(self):
+        from numpy.core.multiarray import as_index_tuple
+
+        def make_badlen_sequence(base):
+            class cls(base):
+                def __len__(self): raise ValueError
+                def __getitem__(self, i): raise IndexError
+            return cls()
+
+        idx = make_badlen_sequence(object)
+        assert_raises(ValueError, len, idx)
+        assert_equal(as_index_tuple(idx), (idx,))
+
+        idx = make_badlen_sequence(tuple)
+        assert_raises(ValueError, len, idx)
+        assert_raises(ValueError, as_index_tuple, idx)
+
+
 class TestFieldIndexing(TestCase):
     def test_scalar_return_type(self):
         # Field access on an array should return an array, even if it

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
@@ -12,7 +12,7 @@
 from . import function_base
 import numpy.matrixlib as matrix
 from .function_base import diff
-from numpy.core.multiarray import ravel_multi_index, unravel_index
+from numpy.core.multiarray import ravel_multi_index, unravel_index, as_index_tuple
 from numpy.lib.stride_tricks import as_strided
 
 makemat = matrix.matrix