Skip to content

Ellipsis handling #168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 68 additions & 3 deletions zarr/tests/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,12 @@ def test_array_1d(self):
assert_array_equal(a[:10], z[:10])
assert_array_equal(a[10:20], z[10:20])
assert_array_equal(a[-10:], z[-10:])
assert_array_equal(a[:10, ...], z[:10, ...])
assert_array_equal(a[10:20, ...], z[10:20, ...])
assert_array_equal(a[-10:, ...], z[-10:, ...])
assert_array_equal(a[..., :10], z[..., :10])
assert_array_equal(a[..., 10:20], z[..., 10:20])
assert_array_equal(a[..., -10:], z[..., -10:])
# ...across chunk boundaries...
assert_array_equal(a[:110], z[:110])
assert_array_equal(a[190:310], z[190:310])
Expand All @@ -135,6 +141,18 @@ def test_array_1d(self):
eq(a[42], z[np.int32(42)])
eq(a[42], z[np.uint64(42)])
eq(a[42], z[np.uint32(42)])
# too many indices
with assert_raises(IndexError):
z[:, :]
with assert_raises(IndexError):
z[0, :]
with assert_raises(IndexError):
z[:, 0]
with assert_raises(IndexError):
z[0, 0]
# only single ellipsis allowed
with assert_raises(IndexError):
z[..., ...]

# check partial assignment
b = np.arange(1e5, 2e5)
Expand Down Expand Up @@ -194,37 +212,84 @@ def test_array_2d(self):
eq(a.nbytes, z.nbytes)
eq(50, z.nchunks_initialized)

# check slicing
# check array-like
assert_array_equal(a, np.array(z))

# check slicing

# total slice
assert_array_equal(a, z[:])
assert_array_equal(a, z[...])
# noinspection PyTypeChecker
assert_array_equal(a, z[slice(None)])

# slice first dimension
assert_array_equal(a[:10], z[:10])
assert_array_equal(a[10:20], z[10:20])
assert_array_equal(a[-10:], z[-10:])
assert_array_equal(a[:10, :], z[:10, :])
assert_array_equal(a[10:20, :], z[10:20, :])
assert_array_equal(a[-10:, :], z[-10:, :])
assert_array_equal(a[:10, ...], z[:10, ...])
assert_array_equal(a[10:20, ...], z[10:20, ...])
assert_array_equal(a[-10:, ...], z[-10:, ...])
assert_array_equal(a[:10, :, ...], z[:10, :, ...])
assert_array_equal(a[10:20, :, ...], z[10:20, :, ...])
assert_array_equal(a[-10:, :, ...], z[-10:, :, ...])

# slice second dimension
assert_array_equal(a[:, :2], z[:, :2])
assert_array_equal(a[:, 2:4], z[:, 2:4])
assert_array_equal(a[:, -2:], z[:, -2:])
assert_array_equal(a[..., :2], z[..., :2])
assert_array_equal(a[..., 2:4], z[..., 2:4])
assert_array_equal(a[..., -2:], z[..., -2:])
assert_array_equal(a[:, ..., :2], z[:, ..., :2])
assert_array_equal(a[:, ..., 2:4], z[:, ..., 2:4])
assert_array_equal(a[:, ..., -2:], z[:, ..., -2:])

# slice both dimensions
assert_array_equal(a[:10, :2], z[:10, :2])
assert_array_equal(a[10:20, 2:4], z[10:20, 2:4])
assert_array_equal(a[-10:, -2:], z[-10:, -2:])
# ...across chunk boundaries...

# slicing across chunk boundaries
assert_array_equal(a[:110], z[:110])
assert_array_equal(a[190:310], z[190:310])
assert_array_equal(a[-110:], z[-110:])
assert_array_equal(a[:110, :], z[:110, :])
assert_array_equal(a[190:310, :], z[190:310, :])
assert_array_equal(a[-110:, :], z[-110:, :])
assert_array_equal(a[:, :3], z[:, :3])
assert_array_equal(a[:, 3:7], z[:, 3:7])
assert_array_equal(a[:, -3:], z[:, -3:])
assert_array_equal(a[:110, :3], z[:110, :3])
assert_array_equal(a[190:310, 3:7], z[190:310, 3:7])
assert_array_equal(a[-110:, -3:], z[-110:, -3:])
# single item

# single row/col/item
assert_array_equal(a[0], z[0])
assert_array_equal(a[-1], z[-1])
assert_array_equal(a[:, 0], z[:, 0])
assert_array_equal(a[:, -1], z[:, -1])
eq(a[0, 0], z[0, 0])
eq(a[-1, -1], z[-1, -1])

# too many indices
with assert_raises(IndexError):
z[:, :, :]
with assert_raises(IndexError):
z[0, :, :]
with assert_raises(IndexError):
z[:, 0, :]
with assert_raises(IndexError):
z[:, :, 0]
with assert_raises(IndexError):
z[0, 0, 0]
# only single ellipsis allowed
with assert_raises(IndexError):
z[..., ...]

# check partial assignment
b = np.arange(10000, 20000).reshape((1000, 10))
z[190:310, 3:7] = b[190:310, 3:7]
Expand Down
21 changes: 15 additions & 6 deletions zarr/tests/test_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,18 +111,17 @@ def test_normalize_array_selection():
eq((slice(0, 100),), normalize_array_selection(slice(None), (100,)))
eq((slice(0, 100),), normalize_array_selection(slice(None, 100), (100,)))
eq((slice(0, 100),), normalize_array_selection(slice(0, None), (100,)))
eq((slice(0, 100),), normalize_array_selection((slice(None), Ellipsis), (100,)))
eq((slice(0, 100),), normalize_array_selection((Ellipsis, slice(None)), (100,)))

# 2D, single item
eq((0, 0), normalize_array_selection((0, 0), (100, 100)))
eq((99, 1), normalize_array_selection((-1, 1), (100, 100)))

# 2D, single col/row
eq((0, slice(0, 100)), normalize_array_selection((0, slice(None)),
(100, 100)))
eq((0, slice(0, 100)), normalize_array_selection((0,),
(100, 100)))
eq((slice(0, 100), 0), normalize_array_selection((slice(None), 0),
(100, 100)))
eq((0, slice(0, 100)), normalize_array_selection((0, slice(None)), (100, 100)))
eq((0, slice(0, 100)), normalize_array_selection((0,), (100, 100)))
eq((slice(0, 100), 0), normalize_array_selection((slice(None), 0), (100, 100)))

# 2D slice
eq((slice(0, 100), slice(0, 100)),
Expand All @@ -131,6 +130,16 @@ def test_normalize_array_selection():
normalize_array_selection(slice(None), (100, 100)))
eq((slice(0, 100), slice(0, 100)),
normalize_array_selection((slice(None), slice(None)), (100, 100)))
eq((slice(0, 100), slice(0, 100)),
normalize_array_selection((Ellipsis, slice(None)), (100, 100)))
eq((slice(0, 100), slice(0, 100)),
normalize_array_selection((slice(None), Ellipsis), (100, 100)))
eq((slice(0, 100), slice(0, 100)),
normalize_array_selection((slice(None), Ellipsis, slice(None)), (100, 100)))
eq((slice(0, 100), slice(0, 100)),
normalize_array_selection((Ellipsis, slice(None), slice(None)), (100, 100)))
eq((slice(0, 100), slice(0, 100)),
normalize_array_selection((slice(None), slice(None), Ellipsis), (100, 100)))

with assert_raises(TypeError):
normalize_array_selection('foo', (100,))
Expand Down
53 changes: 32 additions & 21 deletions zarr/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,29 +192,40 @@ def normalize_array_selection(item, shape):
"""Convenience function to normalize a selection within an array with
the given `shape`."""

# normalize item
if isinstance(item, numbers.Integral):
item = (int(item),)
elif isinstance(item, slice):
# ensure tuple
if not isinstance(item, tuple):
item = (item,)
elif item == Ellipsis:
item = (slice(None),)

# handle tuple of indices/slices
if isinstance(item, tuple):

# determine start and stop indices for all axes
selection = tuple(normalize_axis_selection(i, l)
for i, l in zip(item, shape))

# fill out selection if not completely specified
if len(selection) < len(shape):
selection += tuple(slice(0, l) for l in shape[len(selection):])

return selection

else:
raise TypeError('expected indices or slice, found: %r' % item)
# handle ellipsis
n_ellipsis = sum(1 for i in item if i == Ellipsis)
if n_ellipsis > 1:
raise IndexError("an index can only have a single ellipsis ('...')")
elif n_ellipsis == 1:
n_items_l = item.index(Ellipsis) # items to left of ellipsis
n_items_r = len(item) - (n_items_l + 1) # items to right of ellipsis
n_items = len(item) - 1 # all non-ellipsis items
if n_items >= len(shape):
# ellipsis does nothing, just remove it
item = tuple(i for i in item if i != Ellipsis)
else:
# replace ellipsis with as many slices are needed for number of dims
new_item = item[:n_items_l] + ((slice(None),) * (len(shape) - n_items))
if n_items_r:
new_item += item[-n_items_r:]
item = new_item

# check dimensionality
if len(item) > len(shape):
raise IndexError('too many indices for array')

# determine start and stop indices for all axes
selection = tuple(normalize_axis_selection(i, l) for i, l in zip(item, shape))

# fill out selection if not completely specified
if len(selection) < len(shape):
selection += tuple(slice(0, l) for l in shape[len(selection):])

return selection


def get_chunk_range(selection, chunks):
Expand Down