From 6f4d2478aa32725ef24c57b19150de8e3761f925 Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Mon, 13 Jul 2026 13:57:48 +0300 Subject: [PATCH] gh-58338: Support multi-dimensional slicing in memoryview Slicing a multi-dimensional memoryview, or indexing it with a combination of integers, slices and an ellipsis, now returns a sub-view instead of raising NotImplementedError. Each integer drops the corresponding dimension, a slice keeps it, and the ellipsis expands to full slices over the remaining dimensions. Such sub-views also support slice assignment, with the same shape and format matching rules as one-dimensional assignment. Buffers with suboffsets (indirect, PIL-style) are not yet handled on the sub-view paths and still raise NotImplementedError. --- Doc/library/stdtypes.rst | 41 +++ Doc/whatsnew/3.16.rst | 6 + Lib/test/test_buffer.py | 52 +++- Lib/test/test_memoryview.py | 187 ++++++++++++- ...6-07-13-11-20-00.gh-issue-58338.Kp3xQ7.rst | 4 + Objects/memoryobject.c | 257 ++++++++++++++---- 6 files changed, 472 insertions(+), 75 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-07-13-11-20-00.gh-issue-58338.Kp3xQ7.rst diff --git a/Doc/library/stdtypes.rst b/Doc/library/stdtypes.rst index 886648e820f071d..9ec315700e202cc 100644 --- a/Doc/library/stdtypes.rst +++ b/Doc/library/stdtypes.rst @@ -4660,6 +4660,27 @@ copying. >>> m[::2].tolist() [-11111111, -33333333] + Slicing a multi-dimensional memoryview, + or indexing it with fewer integers than it has dimensions, also results in a sub-view. + A key may mix slices, integers and up to one ellipsis; + each integer drops the corresponding dimension, + and the ellipsis expands to full slices over the remaining dimensions. + An empty tuple or a bare ellipsis therefore selects the whole view:: + + >>> m = memoryview(bytearray(range(12))).cast('B', shape=[3, 4]) + >>> m.tolist() + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]] + >>> m[1].tolist() # partial index: a row sub-view + [4, 5, 6, 7] + >>> m[1:3, 1:3].tolist() # rectangular sub-view + [[5, 6], [9, 10]] + >>> m[1, 1:3].tolist() # mix an index and a slice + [5, 6] + >>> m[..., 1].tolist() # ellipsis fills the leading dimensions + [1, 5, 9] + >>> m[...].tolist() # the whole view + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]] + If the underlying object is writable, the memoryview supports one-dimensional slice assignment. Resizing is not allowed:: @@ -4681,6 +4702,20 @@ copying. >>> data bytearray(b'z1spam') + Multi-dimensional views and sub-views support slice assignment as well, + provided the right-hand side is a :term:`bytes-like object` + whose shape and format match the selected sub-view. + A one-dimensional sub-view accepts any matching bytes-like object, + while a multi-dimensional one needs an object of the same shape:: + + >>> m = memoryview(bytearray(range(12))).cast('B', shape=[3, 4]) + >>> m[0] = b'ABCD' # a one-dimensional row sub-view + >>> m.tolist() + [[65, 66, 67, 68], [4, 5, 6, 7], [8, 9, 10, 11]] + >>> m[1:3, 1:3] = memoryview(bytes([10, 11, 12, 13])).cast('B', shape=[2, 2]) + >>> m.tolist() + [[65, 66, 67, 68], [4, 10, 11, 7], [8, 12, 13, 11]] + One-dimensional memoryviews of :term:`hashable` (read-only) types with formats 'B', 'b' or 'c' are also hashable. The hash is defined as ``hash(m) == hash(m.tobytes())``:: @@ -4707,6 +4742,12 @@ copying. .. versionchanged:: 3.14 memoryview is now a :term:`generic type`. + .. versionchanged:: next + Slicing a multi-dimensional memoryview, + or indexing it with a combination of integers, slices and an ellipsis, + now returns a sub-view, and such sub-views support slice assignment. + Previously these operations raised :exc:`NotImplementedError`. + :class:`memoryview` has several methods: .. method:: __eq__(exporter) diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index da98f1ad6b9bc3b..29c60b060e444de 100644 --- a/Doc/whatsnew/3.16.rst +++ b/Doc/whatsnew/3.16.rst @@ -75,6 +75,12 @@ New features Other language changes ====================== +* Slicing a multi-dimensional :class:`memoryview`, or indexing it with a + combination of integers, slices and an ellipsis, now returns a sub-view, + and such sub-views support slice assignment. + Previously these operations raised :exc:`NotImplementedError`. + (Contributed by Serhiy Storchaka in :gh:`58338`.) + New modules diff --git a/Lib/test/test_buffer.py b/Lib/test/test_buffer.py index 3213a4751273431..73831328c51dcb3 100644 --- a/Lib/test/test_buffer.py +++ b/Lib/test/test_buffer.py @@ -1680,8 +1680,8 @@ def test_ndarray_slice_invalid(self): self.assertRaises(TypeError, nd.__getitem__, ("@%$", slice(0,1,1))) self.assertRaises(TypeError, nd.__getitem__, (slice(0,1,1), {})) - # memoryview: not implemented - self.assertRaises(NotImplementedError, mv.__getitem__, + # memoryview: multi-dimensional slicing validates each slice + self.assertRaises(ValueError, mv.__getitem__, (slice(0,1,1), slice(0,1,0))) self.assertRaises(TypeError, mv.__getitem__, "@%$") @@ -1712,7 +1712,7 @@ def test_ndarray_slice_invalid(self): mr = memoryview(xr) self.assertRaises(ValueError, xl.__setitem__, slice(0,1,1), xr[7:8]) self.assertEqual(xl.tolist(), [[1,2,3,4], [5,6,7,8]]) - self.assertRaises(NotImplementedError, ml.__setitem__, slice(0,1,1), + self.assertRaises(ValueError, ml.__setitem__, slice(0,1,1), mr[7:8]) # differing shape @@ -1789,6 +1789,22 @@ def test_ndarray_slice_multidim(self): else: self.assertEqual(ndsliced.tolist(), sliced) + def test_memoryview_suboffsets_subview(self): + # Sub-views of a buffer with suboffsets (PIL-style indirect buffer) + # are not implemented (gh-58338). A full index and a leading-dimension + # slice follow the suboffsets and still work. + nd = ndarray(list(range(12)), shape=[3, 4], flags=ND_PIL|ND_WRITABLE) + m = memoryview(nd) + self.assertIsNotNone(m.suboffsets) + self.assertEqual(m[1, 2], 6) + self.assertEqual(m[1:3].tolist(), [[4, 5, 6, 7], [8, 9, 10, 11]]) + self.assertRaises(NotImplementedError, m.__getitem__, 1) + self.assertRaises(NotImplementedError, m.__getitem__, + (slice(1, 3), slice(1, 3))) + self.assertRaises(NotImplementedError, m.__getitem__, (Ellipsis, 1)) + self.assertRaises(NotImplementedError, m.__setitem__, 1, + b'\x00\x00\x00\x00') + def test_ndarray_slice_redundant_suboffsets(self): shape_t = (2, 3, 5, 2) ndim = len(shape_t) @@ -2982,9 +2998,15 @@ def test_memoryview_index(self): self.assertRaises(TypeError, m.__getitem__, (0, 0, 0)) self.assertRaises(TypeError, m.__getitem__, (0.0, 0.0)) - # Not implemented: multidimensional sub-views - self.assertRaises(NotImplementedError, m.__getitem__, ()) - self.assertRaises(NotImplementedError, m.__getitem__, 0) + # Empty-tuple or bare-ellipsis indexing of an N-D view returns the + # whole view as a sub-view. + self.assertEqual(m[()].tolist(), ex.tolist()) + self.assertEqual(m[...].tolist(), ex.tolist()) + + # Integer indexing of an N-D view returns a sub-view. + self.assertEqual(m[0].tolist(), [0, 1, 2, 3]) + self.assertEqual(m[-1].tolist(), [8, 9, 10, 11]) + self.assertEqual([row.tolist() for row in m], ex.tolist()) def test_memoryview_assign(self): @@ -3092,8 +3114,11 @@ def test_memoryview_assign(self): self.assertRaises(TypeError, m.__setitem__, (0, 0, 0), 0) self.assertRaises(TypeError, m.__setitem__, (0.0, 0.0), 0) - # Not implemented: multidimensional sub-views - self.assertRaises(NotImplementedError, m.__setitem__, 0, [2, 3]) + # Sub-view assignment needs a bytes-like rvalue with a matching + # shape; a list is not a buffer. + self.assertRaises(TypeError, m.__setitem__, 0, [2, 3]) + m[0] = m[1] + self.assertEqual(m[0].tolist(), m[1].tolist()) def test_memoryview_slice(self): @@ -3105,16 +3130,19 @@ def test_memoryview_slice(self): self.assertRaises(ValueError, m.__setitem__, slice(0,2,0), bytearray([1,2])) - # 0-dim slicing (identity function) - self.assertRaises(NotImplementedError, m.__getitem__, ()) + # empty tuple / bare ellipsis select the whole view + self.assertEqual(m[()].tolist(), list(range(12))) + self.assertEqual(m[...].tolist(), list(range(12))) # multidimensional slices ex = ndarray(list(range(12)), shape=[12], flags=ND_WRITABLE) m = memoryview(ex) - self.assertRaises(NotImplementedError, m.__getitem__, + # too many indices for a 1-D view + self.assertRaises(TypeError, m.__getitem__, (slice(0,2,1), slice(0,2,1))) - self.assertRaises(NotImplementedError, m.__setitem__, + # too many indices for a 1-D view (assignment) + self.assertRaises(TypeError, m.__setitem__, (slice(0,2,1), slice(0,2,1)), bytearray([1,2])) # invalid slice tuple diff --git a/Lib/test/test_memoryview.py b/Lib/test/test_memoryview.py index df5ab9f2c879ca3..95c9f1d951ae221 100644 --- a/Lib/test/test_memoryview.py +++ b/Lib/test/test_memoryview.py @@ -160,9 +160,9 @@ def setitem(key, value): self.assertRaises(TypeError, setitem, (0, slice(0,1,1)), b"a") self.assertRaises(TypeError, setitem, (0,), b"a") self.assertRaises(TypeError, setitem, "a", b"a") - # Not implemented: multidimensional slices + # Too many indices for a 1-D view. slices = (slice(0,1,1), slice(0,1,2)) - self.assertRaises(NotImplementedError, setitem, slices, b"a") + self.assertRaises(TypeError, setitem, slices, b"a") # Trying to resize the memory object exc = ValueError if m.format == 'c' else TypeError self.assertRaises(exc, setitem, 0, b"") @@ -473,7 +473,8 @@ def __len__(self): # Variations on source objects for the buffer: bytes-like objects, then arrays # with itemsize > 1. -# NOTE: support for multi-dimensional objects is unimplemented. +# These sources are all one-dimensional; multi-dimensional views are +# exercised in MultidimTest. class BaseBytesMemoryTests(AbstractMemoryTests): ro_type = bytes @@ -688,6 +689,186 @@ class ArrayMemorySliceSliceTest(unittest.TestCase, pass +class MultidimTest(unittest.TestCase): + # Multi-dimensional slicing and partial indexing return sub-views, and + # such sub-views support assignment (gh-58338). + + def grid(self, shape=(3, 4)): + n = 1 + for d in shape: + n *= d + return memoryview(bytearray(range(n))).cast('B', shape=list(shape)) + + def test_partial_index_subview(self): + m = self.grid() + row = m[1] + self.assertIsInstance(row, memoryview) + self.assertEqual(row.shape, (4,)) + self.assertEqual(row.tolist(), [4, 5, 6, 7]) + self.assertEqual(m[-1].tolist(), [8, 9, 10, 11]) + + def test_multidim_slice_subview(self): + m = self.grid() + sub = m[1:3, 1:3] + self.assertIsInstance(sub, memoryview) + self.assertEqual(sub.shape, (2, 2)) + self.assertEqual(sub.tolist(), [[5, 6], [9, 10]]) + + def test_mixed_index_and_slice(self): + m = self.grid() + self.assertEqual(m[1, 1:3].tolist(), [5, 6]) + self.assertEqual(m[0:2, 2].tolist(), [2, 6]) + + def test_ellipsis(self): + # Ellipsis combined with an integer index (drops a dimension) or a + # slice (keeps it). + m = self.grid() + self.assertEqual(m[..., 1].tolist(), [1, 5, 9]) # column + self.assertEqual(m[1, ...].tolist(), [4, 5, 6, 7]) # row + self.assertEqual(m[..., 1:3].tolist(), [[1, 2], [5, 6], [9, 10]]) + self.assertEqual(m[0:2, ...].tolist(), [[0, 1, 2, 3], [4, 5, 6, 7]]) + + m = self.grid((2, 3, 4)) + self.assertEqual(m[0, ..., 1].tolist(), [1, 5, 9]) + self.assertEqual(m[..., 1].shape, (2, 3)) + + def test_full_index_returns_element(self): + m = self.grid() + self.assertEqual(m[1, 2], 6) + + def test_whole_view(self): + # The empty tuple and a bare ellipsis both select the whole view. + m = self.grid() + for key in ((), ...): + with self.subTest(key=key): + sub = m[key] + self.assertIsInstance(sub, memoryview) + self.assertIsNot(sub, m) + self.assertEqual(sub.shape, (3, 4)) + self.assertEqual(sub.tolist(), m.tolist()) + + def test_whole_view_assign(self): + m = self.grid() + m[...] = memoryview(bytes(range(100, 112))).cast('B', shape=[3, 4]) + self.assertEqual(m[0].tolist(), [100, 101, 102, 103]) + m[()] = memoryview(bytes(range(12))).cast('B', shape=[3, 4]) + self.assertEqual(m[0].tolist(), [0, 1, 2, 3]) + + def test_ellipsis_assign(self): + # Ellipsis combined with an integer index (drops a dimension) or a + # slice (keeps it) on the write path. + m = self.grid() + m[..., 1] = memoryview(bytes([90, 91, 92])) # column, strided + self.assertEqual([m[i, 1] for i in range(3)], [90, 91, 92]) + + m = self.grid() + m[1, ...] = memoryview(bytes([90, 91, 92, 93])) # row + self.assertEqual(m[1].tolist(), [90, 91, 92, 93]) + + m = self.grid() + m[..., 1:3] = memoryview(bytes(range(90, 96))).cast('B', shape=[3, 2]) + self.assertEqual(m.tolist(), + [[0, 90, 91, 3], [4, 92, 93, 7], [8, 94, 95, 11]]) + + m = self.grid((2, 3, 4)) + m[0, ..., 1] = memoryview(bytes([90, 91, 92])) + self.assertEqual([m[0, j, 1] for j in range(3)], [90, 91, 92]) + + def test_step_slice(self): + m = self.grid((4, 4)) + self.assertEqual(m[::2, ::2].tolist(), [[0, 2], [8, 10]]) + self.assertEqual(m[::-1].tolist(), + [[12, 13, 14, 15], [8, 9, 10, 11], + [4, 5, 6, 7], [0, 1, 2, 3]]) + + def test_step_slice_assign(self): + m = self.grid() + m[::2, ::2] = memoryview(bytes([90, 91, 92, 93])).cast('B', shape=[2, 2]) + self.assertEqual(m.tolist(), + [[90, 1, 91, 3], [4, 5, 6, 7], [92, 9, 93, 11]]) + + def test_empty_slice(self): + m = self.grid() + self.assertEqual(m[1:1].shape, (0, 4)) + self.assertEqual(m[1:1].tolist(), []) + self.assertEqual(m[0:0, :].shape, (0, 4)) + + def test_index_out_of_bounds(self): + m = self.grid() + for key in (3, -4, (0, 4), (0, -5), (3, 0)): + with self.subTest(key=key): + self.assertRaises(IndexError, m.__getitem__, key) + + def test_too_many_indices(self): + m = self.grid() + self.assertRaises(TypeError, m.__getitem__, (0, 0, 0)) + + def test_nonbyte_format(self): + # itemsize > 1: strides and offsets are in bytes, so sub-views must + # scale correctly by the element size. + buf = bytearray(struct.pack('12i', *range(12))) + m = memoryview(buf).cast('i', shape=[3, 4]) + self.assertEqual(m.itemsize, struct.calcsize('i')) + self.assertEqual(m[1].tolist(), [4, 5, 6, 7]) + self.assertEqual(m[1, 2], 6) + self.assertEqual(m[::2, ::2].tolist(), [[0, 2], [8, 10]]) + # the rvalue must share the 'i' format, not just the bytes + m[1] = memoryview(struct.pack('4i', 40, 41, 42, 43)).cast('i') + self.assertEqual(m[1].tolist(), [40, 41, 42, 43]) + self.assertRaises(ValueError, m.__setitem__, 1, + struct.pack('4i', 0, 0, 0, 0)) + + def test_assign_block_shape_mismatch(self): + m = self.grid() + self.assertRaises(ValueError, m.__setitem__, (slice(1, 3), slice(1, 3)), + memoryview(bytes([1, 2, 3, 4]))) + + def test_iter_yields_subviews(self): + m = self.grid() + self.assertEqual([row.tolist() for row in m], + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]) + + def test_three_dimensions(self): + m = self.grid((2, 3, 4)) + self.assertEqual(m[1].shape, (3, 4)) + self.assertEqual(m[1, 2].tolist(), [20, 21, 22, 23]) + self.assertEqual(m[1, 2, 3], 23) # full index: a single element + self.assertEqual(m[0, 1:3, 1:3].tolist(), [[5, 6], [9, 10]]) + + def test_assign_row(self): + m = self.grid() + m[0] = bytes([100, 101, 102, 103]) + self.assertEqual(m[0].tolist(), [100, 101, 102, 103]) + + def test_assign_block(self): + m = self.grid() + m[1:3, 1:3] = memoryview(bytes([10, 11, 12, 13])).cast('B', shape=[2, 2]) + self.assertEqual(m.tolist(), + [[0, 1, 2, 3], [4, 10, 11, 7], [8, 12, 13, 11]]) + + def test_assign_element(self): + m = self.grid() + m[1, 2] = 99 + self.assertEqual(m[1, 2], 99) + + def test_assign_from_subview(self): + m = self.grid() + m[0] = m[2] + self.assertEqual(m[0].tolist(), m[2].tolist()) + + def test_assign_shape_mismatch(self): + m = self.grid() + self.assertRaises(ValueError, m.__setitem__, 0, bytes([1, 2])) + + def test_assign_non_buffer(self): + m = self.grid() + self.assertRaises(TypeError, m.__setitem__, 0, [1, 2, 3, 4]) + + def test_assign_readonly(self): + m = memoryview(bytes(range(12))).cast('B', shape=[3, 4]) + self.assertRaises(TypeError, m.__setitem__, 0, bytes([9, 9, 9, 9])) + + class OtherTest(unittest.TestCase): def test_ctypes_cast(self): # Issue 15944: Allow all source formats when casting to bytes. diff --git a/Misc/NEWS.d/next/Library/2026-07-13-11-20-00.gh-issue-58338.Kp3xQ7.rst b/Misc/NEWS.d/next/Library/2026-07-13-11-20-00.gh-issue-58338.Kp3xQ7.rst new file mode 100644 index 000000000000000..5ac3cd47660d9d5 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-13-11-20-00.gh-issue-58338.Kp3xQ7.rst @@ -0,0 +1,4 @@ +Slicing a multi-dimensional :class:`memoryview`, or indexing it with a +combination of integers, slices and an ellipsis, now returns a sub-view, +and such sub-views support slice assignment. +Previously these operations raised :exc:`NotImplementedError`. diff --git a/Objects/memoryobject.c b/Objects/memoryobject.c index d92c7daff15dbf3..ca63ac9e74ef34b 100644 --- a/Objects/memoryobject.c +++ b/Objects/memoryobject.c @@ -2527,6 +2527,36 @@ ptr_from_tuple(const Py_buffer *view, PyObject *tup) return ptr; } +/* Delete dimension 'dim' from a view, shifting the higher dimensions down. + The caller must have already advanced view->buf past that dimension. */ +static void +drop_dimension(Py_buffer *view, int dim) +{ + for (int n = dim; n < view->ndim - 1; n++) { + view->shape[n] = view->shape[n+1]; + view->strides[n] = view->strides[n+1]; + } + view->ndim--; +} + +/* Select integer 'index' along dimension 'dim', advancing buf and dropping + the dimension. Returns -1 with IndexError set if 'index' is out of range. */ +static int +index_dimension(Py_buffer *view, int dim, Py_ssize_t index) +{ + Py_ssize_t nitems = view->shape[dim]; + if (index < 0) + index += nitems; + if (index < 0 || index >= nitems) { + PyErr_Format(PyExc_IndexError, + "index out of bounds on dimension %d", dim + 1); + return -1; + } + view->buf = (char *)view->buf + view->strides[dim] * index; + drop_dimension(view, dim); + return 0; +} + /* Return the item at index. In a one-dimensional view, this is an object with the type specified by view->format. Otherwise, the item is a sub-view. The function is used in memory_subscript() and memory_as_sequence. */ @@ -2554,9 +2584,30 @@ memory_item(PyObject *_self, Py_ssize_t index) return unpack_single(self, ptr, fmt); } - PyErr_SetString(PyExc_NotImplementedError, - "multi-dimensional sub-views are not implemented"); - return NULL; + /* ndim > 1: indexing drops the leading dimension and returns a sub-view. + Buffers exposing suboffsets (PIL-style indirect buffers) are not + handled here yet. */ + if (view->suboffsets != NULL) { + PyErr_SetString(PyExc_NotImplementedError, + "memoryview: sub-views of a buffer with suboffsets are " + "not implemented"); + return NULL; + } + + CHECK_RESTRICTED(self); + + PyMemoryViewObject *sub = + (PyMemoryViewObject *)mbuf_add_view(self->mbuf, view); + if (sub == NULL) + return NULL; + if (index_dimension(&sub->view, 0, index) < 0) { + Py_DECREF(sub); + return NULL; + } + init_len(&sub->view); + init_flags(sub); + + return (PyObject *)sub; } /* Return the item at position *key* (a tuple of indices). */ @@ -2565,7 +2616,6 @@ memory_item_multi(PyMemoryViewObject *self, PyObject *tup) { Py_buffer *view = &(self->view); const char *fmt; - Py_ssize_t nindices = PyTuple_GET_SIZE(tup); char *ptr; CHECK_RELEASED(self); @@ -2574,11 +2624,9 @@ memory_item_multi(PyMemoryViewObject *self, PyObject *tup) if (fmt == NULL) return NULL; - if (nindices < view->ndim) { - PyErr_SetString(PyExc_NotImplementedError, - "sub-views are not implemented"); - return NULL; - } + /* Callers route every key other than a full index to memory_multislice(); + this path only handles a tuple of exactly ndim integer indices. */ + assert(PyTuple_GET_SIZE(tup) == view->ndim); ptr = ptr_from_tuple(view, tup); if (ptr == NULL) return NULL; @@ -2614,40 +2662,109 @@ init_slice(Py_buffer *base, PyObject *key, int dim) return 0; } +/* True if 'key' fully indexes an ndim-dimensional view: a tuple of exactly + ndim integer indices. Such a key selects a single element; every other + tuple, a bare ellipsis, or a shorter/mixed key selects a sub-view. */ static int -is_multislice(PyObject *key) +is_full_index(PyObject *key, int ndim) { - Py_ssize_t size, i; - - if (!PyTuple_Check(key)) - return 0; - size = PyTuple_GET_SIZE(key); - if (size == 0) + if (!PyTuple_Check(key) || PyTuple_GET_SIZE(key) != ndim) return 0; - - for (i = 0; i < size; i++) { - PyObject *x = PyTuple_GET_ITEM(key, i); - if (!PySlice_Check(x)) + for (Py_ssize_t i = 0; i < ndim; i++) { + if (!_PyIndex_Check(PyTuple_GET_ITEM(key, i))) return 0; } return 1; } -static Py_ssize_t -is_multiindex(PyObject *key) +/* mv[key] for a 'key' that mixes slices, integer indices and at most one + ellipsis (NumPy-style basic indexing). A slice keeps its dimension, an + integer index drops it, and the ellipsis expands to full slices over the + remaining dimensions. 'key' is a tuple, or a bare ellipsis standing for the + whole view. Returns a new (sub-)view. + + Read path only for now: buffers exposing suboffsets (PIL-style indirect + buffers) are not handled here and raise NotImplementedError. */ +static PyObject * +memory_multislice(PyMemoryViewObject *self, PyObject *key) { - Py_ssize_t size, i; + Py_buffer *view = &self->view; + int is_tuple = PyTuple_Check(key); + Py_ssize_t nkey = is_tuple ? PyTuple_GET_SIZE(key) : 1; + Py_ssize_t nindexer = 0; /* number of slice/index items (not ellipsis) */ + int have_ellipsis = 0; + + for (Py_ssize_t i = 0; i < nkey; i++) { + PyObject *k = is_tuple ? PyTuple_GET_ITEM(key, i) : key; + if (k == Py_Ellipsis) { + if (have_ellipsis) { + PyErr_SetString(PyExc_TypeError, + "memoryview: more than one ellipsis in slice key"); + return NULL; + } + have_ellipsis = 1; + } + else if (PySlice_Check(k) || _PyIndex_Check(k)) { + nindexer++; + } + else { + PyErr_SetString(PyExc_TypeError, "memoryview: invalid slice key"); + return NULL; + } + } + if (nindexer > view->ndim) { + PyErr_SetString(PyExc_TypeError, + "memoryview: too many indices for memory"); + return NULL; + } + if (view->suboffsets != NULL) { + PyErr_SetString(PyExc_NotImplementedError, + "memoryview: multi-dimensional slicing of a buffer with " + "suboffsets is not supported"); + return NULL; + } - if (!PyTuple_Check(key)) - return 0; - size = PyTuple_GET_SIZE(key); - for (i = 0; i < size; i++) { - PyObject *x = PyTuple_GET_ITEM(key, i); - if (!_PyIndex_Check(x)) { - return 0; + CHECK_RESTRICTED(self); + + PyMemoryViewObject *sliced = + (PyMemoryViewObject *)mbuf_add_view(self->mbuf, view); + if (sliced == NULL) + return NULL; + Py_buffer *dest = &sliced->view; + + int dim = 0; + Py_ssize_t n_ellipsis_dims = view->ndim - nindexer; + for (Py_ssize_t i = 0; i < nkey; i++) { + PyObject *k = is_tuple ? PyTuple_GET_ITEM(key, i) : key; + if (k == Py_Ellipsis) { + dim += (int)n_ellipsis_dims; /* leave these dimensions as-is */ + } + else if (PySlice_Check(k)) { + if (init_slice(dest, k, dim) < 0) { + Py_DECREF(sliced); + return NULL; + } + dim++; + } + else { + /* integer index: drop dimension 'dim'. The higher dimensions + shift down into its place, so 'dim' is not incremented. */ + Py_ssize_t index = PyNumber_AsSsize_t(k, PyExc_IndexError); + if (index == -1 && PyErr_Occurred()) { + Py_DECREF(sliced); + return NULL; + } + if (index_dimension(dest, dim, index) < 0) { + Py_DECREF(sliced); + return NULL; + } } } - return 1; + + init_len(dest); + init_flags(sliced); + + return (PyObject *)sliced; } /* mv[obj] returns an object holding the data for one element if obj @@ -2706,19 +2823,52 @@ memory_subscript(PyObject *_self, PyObject *key) return (PyObject *)sliced; } - else if (is_multiindex(key)) { - return memory_item_multi(self, key); - } - else if (is_multislice(key)) { - PyErr_SetString(PyExc_NotImplementedError, - "multi-dimensional slicing is not implemented"); - return NULL; + else if (key == Py_Ellipsis || PyTuple_Check(key)) { + /* A full index (a tuple of exactly ndim integers) returns the single + element: memory_item_multi() unpacks it in place without allocating + a sub-view, and follows suboffsets. Every other tuple, or a bare + ellipsis, produces a sub-view. The 0-dimensional m[()] case is + handled above. */ + if (is_full_index(key, view->ndim)) { + return memory_item_multi(self, key); + } + CHECK_RESTRICTED(self); + return memory_multislice(self, key); } PyErr_SetString(PyExc_TypeError, "memoryview: invalid slice key"); return NULL; } +/* Assign 'value' to the sub-view selected by 'key' (a slice, ellipsis, or a + tuple mixing slices and integer indices). The lvalue sub-view is produced + by memory_subscript(); 'value' must be an exporter with a matching shape + and format, and is copied in via copy_buffer(). */ +static int +memory_ass_view(PyMemoryViewObject *self, PyObject *key, PyObject *value) +{ + PyObject *lvalue = memory_subscript((PyObject *)self, key); + if (lvalue == NULL) + return -1; + if (!PyMemoryView_Check(lvalue)) { + /* A fully-indexed key yields a scalar, not a sub-view; such keys are + handled by the caller and must not reach here. */ + Py_DECREF(lvalue); + PyErr_SetString(PyExc_TypeError, "memoryview: invalid slice key"); + return -1; + } + + Py_buffer src; + if (PyObject_GetBuffer(value, &src, PyBUF_FULL_RO) < 0) { + Py_DECREF(lvalue); + return -1; + } + int ret = copy_buffer(&((PyMemoryViewObject *)lvalue)->view, &src); + PyBuffer_Release(&src); + Py_DECREF(lvalue); + return ret; +} + static int memory_ass_sub(PyObject *_self, PyObject *key, PyObject *value) { @@ -2758,9 +2908,7 @@ memory_ass_sub(PyObject *_self, PyObject *key, PyObject *value) if (_PyIndex_Check(key)) { Py_ssize_t index; if (1 < view->ndim) { - PyErr_SetString(PyExc_NotImplementedError, - "sub-views are not implemented"); - return -1; + return memory_ass_view(self, key, value); } index = PyNumber_AsSsize_t(key, PyExc_IndexError); if (index == -1 && PyErr_Occurred()) @@ -2797,25 +2945,14 @@ memory_ass_sub(PyObject *_self, PyObject *key, PyObject *value) PyBuffer_Release(&src); return ret; } - if (is_multiindex(key)) { - char *ptr; - if (PyTuple_GET_SIZE(key) < view->ndim) { - PyErr_SetString(PyExc_NotImplementedError, - "sub-views are not implemented"); - return -1; - } + if (is_full_index(key, view->ndim)) { ptr = ptr_from_tuple(view, key); if (ptr == NULL) return -1; return pack_single(self, ptr, value, fmt); } - if (PySlice_Check(key) || is_multislice(key)) { - /* Call memory_subscript() to produce a sliced lvalue, then copy - rvalue into lvalue. This is already implemented in _testbuffer.c. */ - PyErr_SetString(PyExc_NotImplementedError, - "memoryview slice assignments are currently restricted " - "to ndim = 1"); - return -1; + if (PySlice_Check(key) || PyTuple_Check(key) || key == Py_Ellipsis) { + return memory_ass_view(self, key, value); } PyErr_SetString(PyExc_TypeError, "memoryview: invalid slice key"); @@ -3628,6 +3765,11 @@ memoryiter_next(PyObject *self) if (it->it_index < it->it_length) { CHECK_RELEASED(seq); Py_buffer *view = &(seq->view); + if (view->ndim != 1) { + /* Iterating an N-D view yields sub-views over the trailing + dimensions. */ + return memory_item((PyObject *)seq, it->it_index++); + } char *ptr = (char *)seq->view.buf; ptr += view->strides[0] * it->it_index++; @@ -3657,11 +3799,6 @@ memory_iter(PyObject *seq) PyErr_SetString(PyExc_TypeError, "invalid indexing of 0-dim memory"); return NULL; } - if (ndims != 1) { - PyErr_SetString(PyExc_NotImplementedError, - "multi-dimensional sub-views are not implemented"); - return NULL; - } const char *fmt = adjust_fmt(&obj->view); if (fmt == NULL) {