Skip to content

Commit d057157

Browse files
authored
normalize the first indexer (#10948)
* normalize the applied slice * tests for `_index_indexer_1d` * normalize all indexer types * also check `normalize_indexer` * a few more test cases * what's new entry [skip-ci] * reword
1 parent 121f266 commit d057157

File tree

3 files changed

+73
-2
lines changed

3 files changed

+73
-2
lines changed

doc/whats-new.rst

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@ Deprecations
3333

3434
Bug Fixes
3535
~~~~~~~~~
36-
36+
- Always normalize slices when indexing ``LazilyIndexedArray`` instances (:issue:`10941`, :pull:`10948`).
37+
By `Justus Magin <https://github.com/keewis>`_.
3738

3839
Documentation
3940
~~~~~~~~~~~~~

xarray/core/indexing.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,17 @@ def slice_slice_by_array(
355355
return new_indexer
356356

357357

358+
def normalize_indexer(indexer, size):
359+
if isinstance(indexer, slice):
360+
return normalize_slice(indexer, size)
361+
elif isinstance(indexer, np.ndarray):
362+
return normalize_array(indexer, size)
363+
else:
364+
if indexer < 0:
365+
return size + indexer
366+
return indexer
367+
368+
358369
def _index_indexer_1d(
359370
old_indexer: OuterIndexerType,
360371
applied_indexer: OuterIndexerType,
@@ -365,7 +376,7 @@ def _index_indexer_1d(
365376
return old_indexer
366377
if is_full_slice(old_indexer):
367378
# shortcut for full slices
368-
return applied_indexer
379+
return normalize_indexer(applied_indexer, size)
369380

370381
indexer: OuterIndexerType
371382
if isinstance(old_indexer, slice):

xarray/tests/test_indexing.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,24 @@ def test_read_only_view(self) -> None:
275275

276276

277277
class TestLazyArray:
278+
@pytest.mark.parametrize(
279+
["indexer", "size", "expected"],
280+
(
281+
(4, 5, 4),
282+
(-1, 3, 2),
283+
(slice(None), 4, slice(0, 4, 1)),
284+
(slice(1, -3), 7, slice(1, 4, 1)),
285+
(np.array([-1, 3, -2]), 5, np.array([4, 3, 3])),
286+
),
287+
)
288+
def normalize_indexer(self, indexer, size, expected):
289+
actual = indexing.normalize_indexer(indexer, size)
290+
291+
if isinstance(expected, np.ndarray):
292+
np.testing.assert_equal(actual, expected)
293+
else:
294+
assert actual == expected
295+
278296
def test_slice_slice(self) -> None:
279297
arr = ReturnItem()
280298
for size in [100, 99]:
@@ -321,6 +339,47 @@ def test_slice_slice_by_array(self, old_slice, array, size):
321339
expected = np.arange(size)[old_slice][array]
322340
assert_array_equal(actual, expected)
323341

342+
@pytest.mark.parametrize(
343+
["old_indexer", "indexer", "size", "expected"],
344+
(
345+
pytest.param(
346+
slice(None), slice(None, 3), 5, slice(0, 3, 1), id="full_slice-slice"
347+
),
348+
pytest.param(
349+
slice(None), np.arange(2, 4), 5, np.arange(2, 4), id="full_slice-array"
350+
),
351+
pytest.param(slice(None), 3, 5, 3, id="full_slice-int"),
352+
pytest.param(
353+
slice(2, 12, 3), slice(1, 3), 16, slice(5, 11, 3), id="slice_step-slice"
354+
),
355+
pytest.param(
356+
slice(2, 12, 3),
357+
np.array([1, 3]),
358+
16,
359+
np.array([5, 11]),
360+
id="slice_step-array",
361+
),
362+
pytest.param(
363+
np.arange(5), slice(1, 3), 7, np.arange(1, 3), id="array-slice"
364+
),
365+
pytest.param(
366+
np.arange(0, 8, 2),
367+
np.arange(1, 3),
368+
9,
369+
np.arange(2, 6, 2),
370+
id="array-array",
371+
),
372+
pytest.param(np.arange(3), 2, 5, 2, id="array-int"),
373+
),
374+
)
375+
def test_index_indexer_1d(self, old_indexer, indexer, size, expected):
376+
actual = indexing._index_indexer_1d(old_indexer, indexer, size)
377+
378+
if isinstance(expected, np.ndarray):
379+
np.testing.assert_equal(actual, expected)
380+
else:
381+
assert actual == expected
382+
324383
def test_lazily_indexed_array(self) -> None:
325384
original = np.random.rand(10, 20, 30)
326385
x = indexing.NumpyIndexingAdapter(original)

0 commit comments

Comments
 (0)