diff --git a/properties/test_encode_decode.py b/properties/test_encode_decode.py index e7eece7e81e..d636333e018 100644 --- a/properties/test_encode_decode.py +++ b/properties/test_encode_decode.py @@ -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 @@ -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() diff --git a/xarray/testing/strategies.py b/xarray/testing/strategies.py index b76733d113f..d184ad35e8d 100644 --- a/xarray/testing/strategies.py +++ b/xarray/testing/strategies.py @@ -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(),