8000 ENH: Expose prepare_index_tuple as index_tricks.as_index_tuple · numpy/numpy@a86cc8c · GitHub
[go: up one dir, main page]
Skip to content

Commit a86cc8c

Browse files
eric-wiesereric-wieser
committed
ENH: Expose prepare_index_tuple as index_tricks.as_index_tuple
Approaches #8275
1 parent 5afeb16 commit a86cc8c

File tree

7 files changed

+154
-1
lines changed

7 files changed

+154
-1
lines changed

‎doc/release/1.13.0-notes.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,11 @@ being iterated over.
159159
For consistency with ``ndarray`` and ``broadcast``, ``d.ndim`` is a shorthand
160160
for ``len(d.shape)``.
161161

162+
``as_index_tuple`` function in ``index_tricks``
163+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
164+
Exposes the internal normalization that happens when indexing with non-tuple
165+
objects to convert them into tuples. Useful for correctly overriding or
166+
wrapping ndarray.__getitem__
162167

163168
Improvements
164169
============

‎numpy/add_newdocs.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5403,6 +5403,22 @@ def luf(lamdaexpr, *args, **kwargs):
54035403
54045404
""")
54055405

5406+
add_newdoc('numpy.core.multiarray', 'as_index_tuple',
5407+
"""
5408+
as_index_tuple(index)
5409+
5410+
Normalizes an index argument, like that passed to `__getitem__`, into a tuple.
5411+
Follows the invariant that `x[index]` is identical to `x[as_index_tuple(index)].
5412+
5413+
Examples
5414+
--------
5415+
>>> as_index_tuple(1)
5416+
(1,)
5417+
>>> as_index_tuple([1, 2, None])
5418+
(1, 2, None)
5419+
>>> as_index_tuple([1, 2, 3])
5420+
([1, 2, 3])
5421+
""")
54065422

54075423
##############################################################################
54085424
#

‎numpy/core/src/multiarray/mapping.c

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,40 @@ prepare_index_tuple(PyObject *index, PyObject **result)
284284
}
285285
}
286286

287+
/**
288+
* Expose prepare_index_tuple to python code
289+
*/
290+
NPY_NO_EXPORT PyObject *
291+
as_index_tuple(PyObject *NPY_UNUSED(self), PyObject *args)
292+
{
293+
PyObject *obj;
294+
PyObject *result;
295+
PyObject *prepared[NPY_MAXDIMS*2];
296+
npy_intp i, n;
297+
298+
if (!PyArg_ParseTuple(args, "O", &obj)) {
299+
return NULL;
300+
}
301+
n = prepare_index_tuple(obj, prepared);
302+
if (n < 0) {
303+
return NULL;
304+
}
305+
306+
result = PyTuple_New(n);
307+
if (result == NULL) {
308+
return NULL;
309+
}
310+
311+
for (i = 0; i < n; i++) {
312+
PyObject *val = prepared[i];
313+
Py_INCREF(val);
314+
PyTuple_SET_ITEM(result, i, val);
315+
}
316+
317+
return result;
318+
}
319+
320+
287321
/**
288322
* Prepare an npy_index_object from the python slicing object.
289323
*

‎numpy/core/src/multiarray/mapping.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@ array_subscript(PyArrayObject *self, PyObject *op);
4747
NPY_NO_EXPORT int
4848
array_assign_item(PyArrayObject *self, Py_ssize_t i, PyObject *v);
4949

50+
NPY_NO_EXPORT PyObject *
51+
as_index_tuple(PyObject *NPY_UNUSED(self), PyObject *args);
52+
5053
/*
5154
* Prototypes for Mapping calls --- not part of the C-API
5255
* because only useful as part of a getitem call.

‎numpy/core/src/multiarray/multiarraymodule.c

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
6060
#include "templ_common.h" /* for npy_mul_with_overflow_intp */
6161
#include "compiled_base.h"
6262
#include "mem_overlap.h"
63+
#include "mapping.h" /* for as_index_tuple */
6364

6465
/* Only here for API compatibility */
6566
NPY_NO_EXPORT PyTypeObject PyBigArray_Type;
@@ -4293,6 +4294,8 @@ static struct PyMethodDef array_module_methods[] = {
42934294
METH_VARARGS | METH_KEYWORDS, NULL},
42944295
{"normalize_axis_index", (PyCFunction)normalize_axis_index,
42954296
METH_VARARGS | METH_KEYWORDS, NULL},
4297+
{"as_index_tuple", (PyCFunction)as_index_tuple,
4298+
METH_VARARGS, NULL},
42964299
{NULL, NULL, 0, NULL} /* sentinel */
42974300
};
42984301

‎numpy/core/tests/test_indexing.py

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -511,6 +511,98 @@ def test_indexing_array_negative_strides(self):
511511
arr[slices] = 10
512512
assert_array_equal(arr, 10.)
513513

514+
def test_as_index_tuple(self):
515+
from numpy.core.multiarray import as_index_tuple
516+
517+
arr = np.array([1])
518+
sl = np.s_[:]
519+
ell = Ellipsis
520+
obj = object()
521+
522+
# scalars are wrapped in a 1-tuple
523+
assert_equal(as_index_tuple(1), (1,))
524+
assert_equal(as_index_tuple(ell), (ell,))
525+
assert_equal(as_index_tuple(None), (None,))
526+
assert_equal(as_index_tuple(sl), (sl,))
527+
assert_equal(as_index_tuple(arr), (arr,))
528+
assert_equal(as_index_tuple(obj), (obj,))
529+
530+
# tuples are untouched
531+
assert_equal(as_index_tuple((1, 2, 3)), (1, 2, 3))
532+
assert_equal(as_index_tuple((1, 2, ell)), (1, 2, ell))
533+
assert_equal(as_index_tuple((1, 2, None)), (1, 2, None))
534+
assert_equal(as_index_tuple((1, 2, sl)), (1, 2, sl))
535+
assert_equal(as_index_tuple((1, 2, arr)), (1, 2, arr))
536+
537+
# sequences of scalars are wrapped
538+
assert_equal(as_index_tuple([1, 2, 3]), ([1, 2, 3],))
539+
540+
# sequences containing slice objects or ellipses are tuple-ified
541+
assert_equal(as_index_tuple([1, 2, ell]), (1, 2, ell))
542+
assert_equal(as_index_tuple([1, 2, None]), (1, 2, None))
543+
assert_equal(as_index_tuple([1, 2, sl]), (1, 2, sl))
544+
assert_equal(as_index_tuple([1, 2, arr]), (1, 2, arr))
545+
546+
# unless they are >= np.MAXDIMS, in which case they are always wrapped
547+
nd = np.MAXDIMS
548+
assert_equal(as_index_tuple(nd * [1]), (nd * [1],))
549+
assert_equal(as_index_tuple(nd * [ell]), (nd * [ell],))
550+
assert_equal(as_index_tuple(nd * [None]), (nd * [None],))
551+
assert_equal(as_index_tuple(nd * [sl]), (nd * [sl],))
552+
assert_equal(as_index_tuple(nd * [arr]), (nd * [arr],))
553+
554+
def test_as_index_tuple_broken_getitem(self):
555+
from numpy.core.multiarray import as_index_tuple
556+
557+
# test sequences with a broken __getitem__
558+
def make_broken_sequence(base, items):
559+
class Broken(base):
560+
def __len__(self):
561+
return len(items)
562+
def __getitem__(self, i):
563+
val = items[i]
564+
if isinstance(val, Exception):
565+
raise val
566+
return val
567+
return Broken()
568+
569+
# error comes first, so just treat as a scalar
570+
idx = make_broken_sequence(object, [1, ValueError(), None])
571+
assert_raises(ValueError, operator.getitem, idx, 1)
572+
assert_equal(as_index_tuple(idx), (idx,))
573+
574+
# none comes first, so commit to making the tuple
575+
idx = make_broken_sequence(object, [1, None, ValueError()])
576+
assert_raises(ValueError, operator.getitem, idx, 2)
577+
assert_raises(ValueError, as_index_tuple, idx)
578+
579+
# tuples subclasses error in both cases
580+
idx = make_broken_sequence(tuple, [1, ValueError(), None])
581+
assert_raises(ValueError, operator.getitem, idx, 1)
582+
assert_raises(ValueError, as_index_tuple, idx)
583+
584+
idx = make_broken_sequence(tuple, [1, None, ValueError()])
585+
assert_raises(ValueError, operator.getitem, idx, 2)
586+
assert_raises(ValueError, as_index_tuple, idx)
587+
588+
def test_as_index_tuple_broken_len(self):
589+
from numpy.core.multiarray import as_index_tuple
590+
591+
def make_badlen_sequence(base):
592+
class cls(base):
593+
def __len__(self): raise ValueError
594+
def __getitem__(self, i): raise IndexError
595+
return cls()
596+
597+
idx = make_badlen_sequence(object)
598+
assert_raises(ValueError, len, idx)
599+
assert_equal(as_index_tuple(idx), (idx,))
600+
601+
idx = make_badlen_sequence(tuple)
602+
assert_raises(ValueError, len, idx)
603+
assert_raises(ValueError, as_index_tuple, idx)
604+
605+
514606
class TestFieldIndexing(TestCase):
515607
def test_scalar_return_type(self):
516608
# Field access on an array should return an array, even if it

‎numpy/lib/index_tricks.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from . import function_base
1313
import numpy.matrixlib as matrix
1414
from .function_base import diff
15-
from numpy.core.multiarray import ravel_multi_index, unravel_index
15+
from numpy.core.multiarray import ravel_multi_index, unravel_index, as_index_tuple
1616
from numpy.lib.stride_tricks import as_strided
1717

1818
makemat = matrix.matrix

0 commit comments

Comments
 (0)
0