From b14f05e0597061d7bfd4a176a879aaa24b986954 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 20 Oct 2024 22:39:04 +0200 Subject: [PATCH 1/6] check the indexing behavior for scalars --- xarray_array_testing/indexing.py | 40 ++++++++++++++++++++++++ xarray_array_testing/tests/test_numpy.py | 5 +++ 2 files changed, 45 insertions(+) create mode 100644 xarray_array_testing/indexing.py diff --git a/xarray_array_testing/indexing.py b/xarray_array_testing/indexing.py new file mode 100644 index 0000000..4c614b0 --- /dev/null +++ b/xarray_array_testing/indexing.py @@ -0,0 +1,40 @@ +from contextlib import nullcontext + +import hypothesis.strategies as st +import xarray.testing.strategies as xrst +from hypothesis import given + +from xarray_array_testing.base import DuckArrayTestMixin + + +@st.composite +def scalar_indexers(draw, sizes): + # TODO: try to define this using builds and flatmap + possible_indexers = { + dim: st.integers(min_value=-size, max_value=size - 1) + for dim, size in sizes.items() + } + indexers = xrst.unique_subset_of(possible_indexers) + return {dim: draw(indexer) for dim, indexer in draw(indexers).items()} + + +class IndexingTests(DuckArrayTestMixin): + @staticmethod + def expected_errors(op, **parameters): + return nullcontext() + + @given(st.data()) + def test_variable_scalar_isel(self, data): + variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn)) + indexers = data.draw(scalar_indexers(sizes=variable.sizes)) + + with self.expected_errors("scalar_isel", variable=variable): + actual = variable.isel(indexers).data + + raw_indexers = { + dim: indexers.get(dim, slice(None)) for dim in variable.dims + } + expected = variable.data[*raw_indexers.values()] + + assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}" + self.assert_equal(actual, expected) diff --git a/xarray_array_testing/tests/test_numpy.py b/xarray_array_testing/tests/test_numpy.py index 2a9d95b..648f69d 100644 --- a/xarray_array_testing/tests/test_numpy.py +++ b/xarray_array_testing/tests/test_numpy.py @@ -5,6 +5,7 @@ from xarray_array_testing.base import DuckArrayTestMixin from xarray_array_testing.creation import CreationTests +from xarray_array_testing.indexing import IndexingTests from xarray_array_testing.reduction import ReductionTests @@ -32,3 +33,7 @@ class TestCreationNumpy(CreationTests, NumpyTestMixin): class TestReductionNumpy(ReductionTests, NumpyTestMixin): pass + + +class TestIndexingNumpy(IndexingTests, NumpyTestMixin): + pass From ac6a7860926108fec561064404c01f53e2c3ef5c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 20 Oct 2024 23:00:07 +0200 Subject: [PATCH 2/6] refactor the indexers strategy --- xarray_array_testing/indexing.py | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/xarray_array_testing/indexing.py b/xarray_array_testing/indexing.py index 4c614b0..5e80ff0 100644 --- a/xarray_array_testing/indexing.py +++ b/xarray_array_testing/indexing.py @@ -7,15 +7,15 @@ from xarray_array_testing.base import DuckArrayTestMixin +def scalar_indexer(size): + return st.integers(min_value=-size, max_value=size - 1) + + @st.composite -def scalar_indexers(draw, sizes): - # TODO: try to define this using builds and flatmap - possible_indexers = { - dim: st.integers(min_value=-size, max_value=size - 1) - for dim, size in sizes.items() - } - indexers = xrst.unique_subset_of(possible_indexers) - return {dim: draw(indexer) for dim, indexer in draw(indexers).items()} +def indexers(draw, sizes, indexer_strategy_fn): + possible_indexers = {dim: indexer_strategy_fn(size) for dim, size in sizes.items()} + indexers = draw(xrst.unique_subset_of(possible_indexers)) + return {dim: draw(indexer) for dim, indexer in indexers.items()} class IndexingTests(DuckArrayTestMixin): @@ -24,16 +24,14 @@ def expected_errors(op, **parameters): return nullcontext() @given(st.data()) - def test_variable_scalar_isel(self, data): + def test_variable_isel_scalars(self, data): variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn)) - indexers = data.draw(scalar_indexers(sizes=variable.sizes)) + idx = data.draw(indexers(variable.sizes, scalar_indexer)) - with self.expected_errors("scalar_isel", variable=variable): - actual = variable.isel(indexers).data + with self.expected_errors("isel_scalars", variable=variable): + actual = variable.isel(idx).data - raw_indexers = { - dim: indexers.get(dim, slice(None)) for dim in variable.dims - } + raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims} expected = variable.data[*raw_indexers.values()] assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}" From 4e9570cdbe28a9ce503bf9a82365c7e29fc2a90f Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 20 Oct 2024 23:01:48 +0200 Subject: [PATCH 3/6] add a test for slices --- xarray_array_testing/indexing.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/xarray_array_testing/indexing.py b/xarray_array_testing/indexing.py index 5e80ff0..44eaff8 100644 --- a/xarray_array_testing/indexing.py +++ b/xarray_array_testing/indexing.py @@ -36,3 +36,17 @@ def test_variable_isel_scalars(self, data): assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}" self.assert_equal(actual, expected) + + @given(st.data()) + def test_variable_isel_slices(self, data): + variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn)) + idx = data.draw(indexers(variable.sizes, st.slices)) + + with self.expected_errors("isel_slices", variable=variable): + actual = variable.isel(idx).data + + raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims} + expected = variable.data[*raw_indexers.values()] + + assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}" + self.assert_equal(actual, expected) From 41c6cd30f14c219ff12cd963f336ca65c4480224 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 20 Oct 2024 23:14:25 +0200 Subject: [PATCH 4/6] comment on future refactoring work --- xarray_array_testing/indexing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray_array_testing/indexing.py b/xarray_array_testing/indexing.py index 44eaff8..0b1e8f2 100644 --- a/xarray_array_testing/indexing.py +++ b/xarray_array_testing/indexing.py @@ -13,6 +13,7 @@ def scalar_indexer(size): @st.composite def indexers(draw, sizes, indexer_strategy_fn): + # TODO: make use of `flatmap` and `builds` instead of `composite` possible_indexers = {dim: indexer_strategy_fn(size) for dim, size in sizes.items()} indexers = draw(xrst.unique_subset_of(possible_indexers)) return {dim: draw(indexer) for dim, indexer in indexers.items()} From 3e555d7af23007b06af780cbcdf3d13d4f55c62c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 29 Oct 2024 12:53:56 +0100 Subject: [PATCH 5/6] configure coverage --- pyproject.toml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 2d387a6..65693d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,3 +63,11 @@ known-third-party = [] [tool.ruff.lint.flake8-tidy-imports] ban-relative-imports = "all" + +[tool.coverage.run] +source = ["xarray_array_testing"] +branch = true + +[tool.coverage.report] +show_missing = true +exclude_lines = ["pragma: no cover", "if TYPE_CHECKING"] From 99f3f3bb08b77160fadfe5cccabfa7ce75b7446e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 29 Oct 2024 20:22:00 +0100 Subject: [PATCH 6/6] refactor to also allow integer arrays --- xarray_array_testing/indexing.py | 69 ++++++++++++++++++++++---------- 1 file changed, 48 insertions(+), 21 deletions(-) diff --git a/xarray_array_testing/indexing.py b/xarray_array_testing/indexing.py index 0b1e8f2..d7d1a02 100644 --- a/xarray_array_testing/indexing.py +++ b/xarray_array_testing/indexing.py @@ -1,5 +1,6 @@ from contextlib import nullcontext +import hypothesis.extra.numpy as npst import hypothesis.strategies as st import xarray.testing.strategies as xrst from hypothesis import given @@ -11,39 +12,65 @@ def scalar_indexer(size): return st.integers(min_value=-size, max_value=size - 1) +def integer_array_indexer(size): + dtypes = npst.integer_dtypes() + + return npst.arrays( + dtypes, size, elements={"min_value": -size, "max_value": size - 1} + ) + + +def indexers(size, indexer_types): + indexer_strategy_fns = { + "scalars": scalar_indexer, + "slices": st.slices, + "integer_arrays": integer_array_indexer, + } + + bad_types = set(indexer_types) - indexer_strategy_fns.keys() + if bad_types: + raise ValueError(f"unknown indexer strategies: {sorted(bad_types)}") + + # use the order of definition to prefer simpler strategies over more complex + # ones + indexer_strategies = [ + strategy_fn(size) + for name, strategy_fn in indexer_strategy_fns.items() + if name in indexer_types + ] + return st.one_of(*indexer_strategies) + + @st.composite -def indexers(draw, sizes, indexer_strategy_fn): +def orthogonal_indexers(draw, sizes, indexer_types): # TODO: make use of `flatmap` and `builds` instead of `composite` - possible_indexers = {dim: indexer_strategy_fn(size) for dim, size in sizes.items()} - indexers = draw(xrst.unique_subset_of(possible_indexers)) - return {dim: draw(indexer) for dim, indexer in indexers.items()} + possible_indexers = { + dim: indexers(size, indexer_types) for dim, size in sizes.items() + } + concrete_indexers = draw(xrst.unique_subset_of(possible_indexers)) + return {dim: draw(indexer) for dim, indexer in concrete_indexers.items()} class IndexingTests(DuckArrayTestMixin): + @property + def orthogonal_indexer_types(self): + return st.sampled_from(["scalars", "slices"]) + @staticmethod def expected_errors(op, **parameters): return nullcontext() @given(st.data()) - def test_variable_isel_scalars(self, data): - variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn)) - idx = data.draw(indexers(variable.sizes, scalar_indexer)) - - with self.expected_errors("isel_scalars", variable=variable): - actual = variable.isel(idx).data - - raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims} - expected = variable.data[*raw_indexers.values()] - - assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}" - self.assert_equal(actual, expected) - - @given(st.data()) - def test_variable_isel_slices(self, data): + def test_variable_isel_orthogonal(self, data): + indexer_types = data.draw( + st.lists(self.orthogonal_indexer_types, min_size=1, unique=True) + ) variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn)) - idx = data.draw(indexers(variable.sizes, st.slices)) + idx = data.draw(orthogonal_indexers(variable.sizes, indexer_types)) - with self.expected_errors("isel_slices", variable=variable): + with self.expected_errors( + "isel_orthogonal", variable=variable, indexer_types=indexer_types + ): actual = variable.isel(idx).data raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims}