Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion properties/test_encode_decode.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
from hypothesis import given

import xarray as xr
from xarray.testing.strategies import variables
from xarray.testing.strategies import cftime_arrays, variables
from xarray.tests import requires_cftime


@pytest.mark.slow
Expand Down Expand Up @@ -43,3 +44,28 @@ def test_CFScaleOffset_coder_roundtrip(original) -> None:
coder = xr.coding.variables.CFScaleOffsetCoder()
roundtripped = coder.decode(coder.encode(original))
xr.testing.assert_identical(original, roundtripped)


@requires_cftime
@given(original_array=cftime_arrays(shapes=npst.array_shapes(max_dims=1)))
def test_CFDatetime_coder_roundtrip_cftime(original_array) -> None:
original = xr.Variable("time", original_array)
coder = xr.coding.times.CFDatetimeCoder(use_cftime=True)
roundtripped = coder.decode(coder.encode(original))
xr.testing.assert_identical(original, roundtripped)


@given(
original_array=npst.arrays(
dtype=npst.datetime64_dtypes(endianness="=", max_period="ns"),
shape=npst.array_shapes(max_dims=1),
)
)
def test_CFDatetime_coder_roundtrip_numpy(original_array) -> None:
original = xr.Variable("time", original_array)
coder = xr.coding.times.CFDatetimeCoder(use_cftime=False)
roundtripped = coder.decode(coder.encode(original))
xr.testing.assert_identical(original, roundtripped)


# datetime_arrays =, shape=npst.array_shapes(min_dims=1, max_dims=1)) | cftime_arrays()
54 changes: 54 additions & 0 deletions xarray/testing/strategies.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,60 @@ def dimension_names(
)


calendars = st.sampled_from(
[
"standard",
"gregorian",
"proleptic_gregorian",
"noleap",
"365_day",
"360_day",
"julian",
"all_leap",
"366_day",
]
)


@st.composite
def cftime_units(draw: st.DrawFn, *, calendar: str) -> str:
choices = ["days", "hours", "minutes", "seconds", "milliseconds", "microseconds"]
if calendar == "360_day":
choices += ["months"]
elif calendar == "noleap":
choices += ["common_years"]
time_units = draw(st.sampled_from(choices))

dt = draw(st.datetimes())
year, month, day = dt.year, dt.month, dt.day
if calendar == "360_day":
day = min(day, 30)
if calendar in ["360_day", "365_day", "noleap"] and month == 2 and day == 29:
day = 28

return f"{time_units} since {year}-{month}-{day}"


@st.composite
def cftime_arrays(
draw: st.DrawFn,
*,
shapes: st.SearchStrategy[tuple[int, ...]] = npst.array_shapes(),
calendars: st.SearchStrategy[str] = calendars,
elements: dict[str, Any] | None = None,
) -> np.ndarray[Any, Any]:
import cftime

if elements is None:
elements = {}
elements.setdefault("min_value", 0)
elements.setdefault("max_value", 10_000)
cal = draw(calendars)
values = draw(npst.arrays(dtype=np.int64, shape=shapes, elements=elements))
unit = draw(cftime_units(calendar=cal))
return cftime.num2date(values, units=unit, calendar=cal)


def dimension_sizes(
*,
dim_names: st.SearchStrategy[Hashable] = names(),
Expand Down