1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
|
import itertools
import warnings
import numpy as np
import pytest
import xarray as xr
from xarray import Dataset
from xarray.testing import _assert_internal_invariants
pytest.importorskip("hypothesis")
pytestmark = pytest.mark.slow_hypothesis
import hypothesis.extra.numpy as npst
import hypothesis.strategies as st
from hypothesis import note, settings
from hypothesis.stateful import (
RuleBasedStateMachine,
initialize,
invariant,
precondition,
rule,
)
import xarray.testing.strategies as xrst
# Strategy for generating names - uniqueness is enforced by the state machine
NAME_STRATEGY = xrst.names()
DIM_NAME = xrst.dimension_names(name_strategy=NAME_STRATEGY, min_dims=1, max_dims=1)
index_variables = st.builds(
xr.Variable,
data=npst.arrays(
dtype=xrst.pandas_index_dtypes(),
shape=npst.array_shapes(min_dims=1, max_dims=1),
elements=dict(allow_nan=False, allow_infinity=False, allow_subnormal=False),
unique=True,
),
dims=DIM_NAME,
attrs=xrst.attrs(),
)
def add_dim_coord_and_data_var(ds, var):
(name,) = var.dims
# dim coord
ds[name] = var
# non-dim coord of same size; this allows renaming
ds[name + "_"] = var
class DatasetStateMachine(RuleBasedStateMachine):
# Can't use bundles because we'd need pre-conditions on consumes(bundle)
# indexed_dims = Bundle("indexed_dims")
# multi_indexed_dims = Bundle("multi_indexed_dims")
def __init__(self):
super().__init__()
self.dataset = Dataset()
self.check_default_indexes = True
# We track these separately as lists so we can guarantee order of iteration over them.
# Order of iteration over Dataset.dims is not guaranteed
self.indexed_dims = []
self.multi_indexed_dims = []
# Track all used names to ensure uniqueness (avoids flaky Hypothesis tests)
self.used_names: set[str] = set()
def _draw_unique_name(self, data) -> str:
"""Draw a name that hasn't been used yet in this test case."""
name = data.draw(NAME_STRATEGY.filter(lambda x: x not in self.used_names))
self.used_names.add(name)
return name
def _draw_unique_var(self, data) -> xr.Variable:
"""Draw an index variable with a unique dimension name."""
var = data.draw(index_variables)
# Replace with a guaranteed unique name
new_name = self._draw_unique_name(data)
return xr.Variable(dims=(new_name,), data=var.data, attrs=var.attrs)
@initialize(data=st.data())
def init_ds(self, data):
"""Initialize the Dataset so that at least one rule will always fire."""
var = self._draw_unique_var(data)
(name,) = var.dims
note(f"initializing with dimension coordinate {name}")
add_dim_coord_and_data_var(self.dataset, var)
self.indexed_dims.append(name)
# TODO: stacking with a timedelta64 index and unstacking converts it to object
@rule(data=st.data())
def add_dim_coord(self, data):
var = self._draw_unique_var(data)
(name,) = var.dims
note(f"adding dimension coordinate {name}")
add_dim_coord_and_data_var(self.dataset, var)
self.indexed_dims.append(name)
@rule(data=st.data())
def assign_coords(self, data):
var = self._draw_unique_var(data)
(name,) = var.dims
note(f"assign_coords: {name}")
self.dataset = self.dataset.assign_coords({name: var})
self.indexed_dims.append(name)
@property
def has_indexed_dims(self) -> bool:
return bool(self.indexed_dims + self.multi_indexed_dims)
@rule(data=st.data())
@precondition(lambda self: self.has_indexed_dims)
def reset_index(self, data):
dim = data.draw(st.sampled_from(self.indexed_dims + self.multi_indexed_dims))
self.check_default_indexes = False
note(f"> resetting {dim}")
self.dataset = self.dataset.reset_index(dim)
if dim in self.indexed_dims:
del self.indexed_dims[self.indexed_dims.index(dim)]
elif dim in self.multi_indexed_dims:
del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]
@rule(data=st.data(), create_index=st.booleans())
@precondition(lambda self: bool(self.indexed_dims))
def stack(self, data, create_index):
newname = self._draw_unique_name(data)
oldnames = data.draw(
st.lists(
st.sampled_from(self.indexed_dims),
min_size=1,
max_size=3 if create_index else None,
unique=True,
)
)
note(f"> stacking {oldnames} as {newname}")
self.dataset = self.dataset.stack(
{newname: oldnames}, create_index=create_index
)
if create_index:
self.multi_indexed_dims += [newname]
# if create_index is False, then we just drop these
for dim in oldnames:
del self.indexed_dims[self.indexed_dims.index(dim)]
@rule(data=st.data())
@precondition(lambda self: bool(self.multi_indexed_dims))
def unstack(self, data):
# TODO: add None
dim = data.draw(st.sampled_from(self.multi_indexed_dims))
note(f"> unstacking {dim}")
if dim is not None:
pd_index = self.dataset.xindexes[dim].index
self.dataset = self.dataset.unstack(dim)
del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]
if dim is not None:
self.indexed_dims.extend(pd_index.names)
else:
# TODO: fix this
pass
@rule(data=st.data())
@precondition(lambda self: bool(self.dataset.variables))
def rename_vars(self, data):
newname = self._draw_unique_name(data)
dim = data.draw(st.sampled_from(sorted(self.dataset.variables)))
# benbovy: "skip the default indexes invariant test when the name of an
# existing dimension coordinate is passed as input kwarg or dict key
# to .rename_vars()."
self.check_default_indexes = False
note(f"> renaming {dim} to {newname}")
self.dataset = self.dataset.rename_vars({dim: newname})
if dim in self.indexed_dims:
del self.indexed_dims[self.indexed_dims.index(dim)]
elif dim in self.multi_indexed_dims:
del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]
@precondition(lambda self: bool(self.dataset.dims))
@rule(data=st.data())
def drop_dims(self, data):
dims = data.draw(
st.lists(
st.sampled_from(sorted(self.dataset.dims)),
min_size=1,
unique=True,
)
)
note(f"> drop_dims: {dims}")
# TODO: dropping a multi-index dimension raises a DeprecationWarning
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=DeprecationWarning)
self.dataset = self.dataset.drop_dims(dims)
for dim in dims:
if dim in self.indexed_dims:
del self.indexed_dims[self.indexed_dims.index(dim)]
elif dim in self.multi_indexed_dims:
del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]
@precondition(lambda self: bool(self.indexed_dims))
@rule(data=st.data())
def drop_indexes(self, data):
self.check_default_indexes = False
dims = data.draw(
st.lists(st.sampled_from(self.indexed_dims), min_size=1, unique=True)
)
note(f"> drop_indexes: {dims}")
self.dataset = self.dataset.drop_indexes(dims)
for dim in dims:
if dim in self.indexed_dims:
del self.indexed_dims[self.indexed_dims.index(dim)]
elif dim in self.multi_indexed_dims:
del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]
@property
def swappable_dims(self):
ds = self.dataset
options = []
for dim in self.indexed_dims:
choices = [
name
for name, var in ds._variables.items()
if var.dims == (dim,)
# TODO: Avoid swapping a dimension to itself
and name != dim
]
options.extend(
(a, b) for a, b in itertools.zip_longest((dim,), choices, fillvalue=dim)
)
return options
@rule(data=st.data())
# TODO: swap_dims is basically all broken if a multiindex is present
# TODO: Avoid swapping from Index to a MultiIndex level
# TODO: Avoid swapping from MultiIndex to a level of the same MultiIndex
# TODO: Avoid swapping when a MultiIndex is present
@precondition(lambda self: not bool(self.multi_indexed_dims))
@precondition(lambda self: bool(self.swappable_dims))
def swap_dims(self, data):
ds = self.dataset
options = self.swappable_dims
dim, to = data.draw(st.sampled_from(options))
note(
f"> swapping {dim} to {to}, found swappable dims: {options}, all_dims: {tuple(self.dataset.dims)}"
)
self.dataset = ds.swap_dims({dim: to})
del self.indexed_dims[self.indexed_dims.index(dim)]
self.indexed_dims += [to]
@invariant()
def assert_invariants(self):
# note(f"> ===\n\n {self.dataset!r} \n===\n\n")
_assert_internal_invariants(self.dataset, self.check_default_indexes)
DatasetStateMachine.TestCase.settings = settings(max_examples=300, deadline=None)
DatasetTest = DatasetStateMachine.TestCase
@pytest.mark.skip(reason="failure detected by hypothesis")
def test_unstack_object():
ds = xr.Dataset()
ds["0"] = np.array(["", "\x000"], dtype=object)
ds.stack({"1": ["0"]}).unstack()
@pytest.mark.skip(reason="failure detected by hypothesis")
def test_unstack_timedelta_index():
ds = xr.Dataset()
ds["0"] = np.array([0, 1, 2, 3], dtype="timedelta64[ns]")
ds.stack({"1": ["0"]}).unstack()
|
