Skip to content

Commit fa5e335

Browse files
authored
vectorized indexing tests (#31)
1 parent c04ea81 commit fa5e335

File tree

1 file changed

+33
-0
lines changed

1 file changed

+33
-0
lines changed

xarray_array_testing/indexing.py

+33
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import hypothesis.extra.numpy as npst
44
import hypothesis.strategies as st
5+
import xarray as xr
56
import xarray.testing.strategies as xrst
67
from hypothesis import given
78

@@ -51,6 +52,24 @@ def orthogonal_indexers(draw, sizes, indexer_types):
5152
return {dim: draw(indexer) for dim, indexer in concrete_indexers.items()}
5253

5354

55+
@st.composite
56+
def vectorized_indexers(draw, sizes):
57+
max_size = max(sizes.values())
58+
shape = draw(st.integers(min_value=1, max_value=max_size))
59+
dtypes = npst.integer_dtypes()
60+
61+
indexers = {
62+
dim: npst.arrays(
63+
dtypes, shape, elements={"min_value": -size, "max_value": size - 1}
64+
)
65+
for dim, size in sizes.items()
66+
}
67+
68+
return {
69+
dim: xr.Variable("points", draw(indexer)) for dim, indexer in indexers.items()
70+
}
71+
72+
5473
class IndexingTests(DuckArrayTestMixin):
5574
@property
5675
def orthogonal_indexer_types(self):
@@ -78,3 +97,17 @@ def test_variable_isel_orthogonal(self, data):
7897

7998
assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}"
8099
self.assert_equal(actual, expected)
100+
101+
@given(st.data())
102+
def test_variable_isel_vectorized(self, data):
103+
variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
104+
idx = data.draw(vectorized_indexers(variable.sizes))
105+
106+
with self.expected_errors("isel_vectorized", variable=variable):
107+
actual = variable.isel(idx).data
108+
109+
raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims}
110+
expected = variable.data[*raw_indexers.values()]
111+
112+
assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}"
113+
self.assert_equal(actual, expected)

0 commit comments

Comments
 (0)