diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 677b2194a55..26c889d241e 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -33,7 +33,8 @@ Deprecations Bug Fixes ~~~~~~~~~ - +- Always normalize slices when indexing ``LazilyIndexedArray`` instances (:issue:`10941`, :pull:`10948`). + By `Justus Magin `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 5d022d0430b..2a86afc0a98 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -355,6 +355,17 @@ def slice_slice_by_array( return new_indexer +def normalize_indexer(indexer, size): + if isinstance(indexer, slice): + return normalize_slice(indexer, size) + elif isinstance(indexer, np.ndarray): + return normalize_array(indexer, size) + else: + if indexer < 0: + return size + indexer + return indexer + + def _index_indexer_1d( old_indexer: OuterIndexerType, applied_indexer: OuterIndexerType, @@ -365,7 +376,7 @@ def _index_indexer_1d( return old_indexer if is_full_slice(old_indexer): # shortcut for full slices - return applied_indexer + return normalize_indexer(applied_indexer, size) indexer: OuterIndexerType if isinstance(old_indexer, slice): diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index ffcfd3116f4..6b564c6f032 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -275,6 +275,24 @@ def test_read_only_view(self) -> None: class TestLazyArray: + @pytest.mark.parametrize( + ["indexer", "size", "expected"], + ( + (4, 5, 4), + (-1, 3, 2), + (slice(None), 4, slice(0, 4, 1)), + (slice(1, -3), 7, slice(1, 4, 1)), + (np.array([-1, 3, -2]), 5, np.array([4, 3, 3])), + ), + ) + def normalize_indexer(self, indexer, size, expected): + actual = indexing.normalize_indexer(indexer, size) + + if isinstance(expected, np.ndarray): + np.testing.assert_equal(actual, expected) + else: + assert actual == expected + def test_slice_slice(self) -> None: arr = ReturnItem() for size in [100, 99]: @@ -321,6 +339,47 @@ def test_slice_slice_by_array(self, old_slice, array, size): expected = np.arange(size)[old_slice][array] assert_array_equal(actual, expected) + @pytest.mark.parametrize( + ["old_indexer", "indexer", "size", "expected"], + ( + pytest.param( + slice(None), slice(None, 3), 5, slice(0, 3, 1), id="full_slice-slice" + ), + pytest.param( + slice(None), np.arange(2, 4), 5, np.arange(2, 4), id="full_slice-array" + ), + pytest.param(slice(None), 3, 5, 3, id="full_slice-int"), + pytest.param( + slice(2, 12, 3), slice(1, 3), 16, slice(5, 11, 3), id="slice_step-slice" + ), + pytest.param( + slice(2, 12, 3), + np.array([1, 3]), + 16, + np.array([5, 11]), + id="slice_step-array", + ), + pytest.param( + np.arange(5), slice(1, 3), 7, np.arange(1, 3), id="array-slice" + ), + pytest.param( + np.arange(0, 8, 2), + np.arange(1, 3), + 9, + np.arange(2, 6, 2), + id="array-array", + ), + pytest.param(np.arange(3), 2, 5, 2, id="array-int"), + ), + ) + def test_index_indexer_1d(self, old_indexer, indexer, size, expected): + actual = indexing._index_indexer_1d(old_indexer, indexer, size) + + if isinstance(expected, np.ndarray): + np.testing.assert_equal(actual, expected) + else: + assert actual == expected + def test_lazily_indexed_array(self) -> None: original = np.random.rand(10, 20, 30) x = indexing.NumpyIndexingAdapter(original)