from __future__ import annotations import pickle from datetime import timedelta from textwrap import dedent import numpy as np import pandas as pd import pytest from packaging.version import Version import xarray as xr from xarray.coding.cftimeindex import ( CFTimeIndex, _parse_array_of_cftime_strings, _parse_iso8601_with_reso, _parsed_string_to_bounds, assert_all_valid_date_type, parse_iso8601_like, ) from xarray.tests import ( assert_array_equal, assert_identical, has_cftime, requires_cftime, ) from xarray.tests.test_coding_times import ( _ALL_CALENDARS, _NON_STANDARD_CALENDARS, _all_cftime_date_types, ) # cftime 1.5.2 renames "gregorian" to "standard" standard_or_gregorian = "" if has_cftime: import cftime if Version(cftime.__version__) >= Version("1.5.2"): standard_or_gregorian = "standard" else: standard_or_gregorian = "gregorian" def date_dict(year=None, month=None, day=None, hour=None, minute=None, second=None): return dict( year=year, month=month, day=day, hour=hour, minute=minute, second=second ) ISO8601_LIKE_STRING_TESTS = { "year": ("1999", date_dict(year="1999")), "month": ("199901", date_dict(year="1999", month="01")), "month-dash": ("1999-01", date_dict(year="1999", month="01")), "day": ("19990101", date_dict(year="1999", month="01", day="01")), "day-dash": ("1999-01-01", date_dict(year="1999", month="01", day="01")), "hour": ("19990101T12", date_dict(year="1999", month="01", day="01", hour="12")), "hour-dash": ( "1999-01-01T12", date_dict(year="1999", month="01", day="01", hour="12"), ), "hour-space-separator": ( "1999-01-01 12", date_dict(year="1999", month="01", day="01", hour="12"), ), "minute": ( "19990101T1234", date_dict(year="1999", month="01", day="01", hour="12", minute="34"), ), "minute-dash": ( "1999-01-01T12:34", date_dict(year="1999", month="01", day="01", hour="12", minute="34"), ), "minute-space-separator": ( "1999-01-01 12:34", date_dict(year="1999", month="01", day="01", hour="12", minute="34"), ), "second": ( "19990101T123456", date_dict( year="1999", month="01", day="01", hour="12", minute="34", second="56" ), ), "second-dash": ( "1999-01-01T12:34:56", date_dict( year="1999", month="01", day="01", hour="12", minute="34", second="56" ), ), "second-space-separator": ( "1999-01-01 12:34:56", date_dict( year="1999", month="01", day="01", hour="12", minute="34", second="56" ), ), } @pytest.mark.parametrize( ("string", "expected"), list(ISO8601_LIKE_STRING_TESTS.values()), ids=list(ISO8601_LIKE_STRING_TESTS.keys()), ) def test_parse_iso8601_like(string, expected): result = parse_iso8601_like(string) assert result == expected with pytest.raises(ValueError): parse_iso8601_like(string + "3") parse_iso8601_like(string + ".3") _CFTIME_CALENDARS = [ "365_day", "360_day", "julian", "all_leap", "366_day", "gregorian", "proleptic_gregorian", ] @pytest.fixture(params=_CFTIME_CALENDARS) def date_type(request): return _all_cftime_date_types()[request.param] @pytest.fixture def index(date_type): dates = [ date_type(1, 1, 1), date_type(1, 2, 1), date_type(2, 1, 1), date_type(2, 2, 1), ] return CFTimeIndex(dates) @pytest.fixture def monotonic_decreasing_index(date_type): dates = [ date_type(2, 2, 1), date_type(2, 1, 1), date_type(1, 2, 1), date_type(1, 1, 1), ] return CFTimeIndex(dates) @pytest.fixture def length_one_index(date_type): dates = [date_type(1, 1, 1)] return CFTimeIndex(dates) @pytest.fixture def da(index): return xr.DataArray([1, 2, 3, 4], coords=[index], dims=["time"]) @pytest.fixture def series(index): return pd.Series([1, 2, 3, 4], index=index) @pytest.fixture def df(index): return pd.DataFrame([1, 2, 3, 4], index=index) @pytest.fixture def feb_days(date_type): import cftime if date_type is cftime.DatetimeAllLeap: return 29 elif date_type is cftime.Datetime360Day: return 30 else: return 28 @pytest.fixture def dec_days(date_type): import cftime if date_type is cftime.Datetime360Day: return 30 else: return 31 @pytest.fixture def index_with_name(date_type): dates = [ date_type(1, 1, 1), date_type(1, 2, 1), date_type(2, 1, 1), date_type(2, 2, 1), ] return CFTimeIndex(dates, name="foo") @requires_cftime @pytest.mark.parametrize(("name", "expected_name"), [("bar", "bar"), (None, "foo")]) def test_constructor_with_name(index_with_name, name, expected_name): result = CFTimeIndex(index_with_name, name=name).name assert result == expected_name @requires_cftime def test_assert_all_valid_date_type(date_type, index): import cftime if date_type is cftime.DatetimeNoLeap: mixed_date_types = np.array( [date_type(1, 1, 1), cftime.DatetimeAllLeap(1, 2, 1)] ) else: mixed_date_types = np.array( [date_type(1, 1, 1), cftime.DatetimeNoLeap(1, 2, 1)] ) with pytest.raises(TypeError): assert_all_valid_date_type(mixed_date_types) with pytest.raises(TypeError): assert_all_valid_date_type(np.array([1, date_type(1, 1, 1)])) assert_all_valid_date_type(np.array([date_type(1, 1, 1), date_type(1, 2, 1)])) @requires_cftime @pytest.mark.parametrize( ("field", "expected"), [ ("year", [1, 1, 2, 2]), ("month", [1, 2, 1, 2]), ("day", [1, 1, 1, 1]), ("hour", [0, 0, 0, 0]), ("minute", [0, 0, 0, 0]), ("second", [0, 0, 0, 0]), ("microsecond", [0, 0, 0, 0]), ], ) def test_cftimeindex_field_accessors(index, field, expected): result = getattr(index, field) assert_array_equal(result, expected) @requires_cftime def test_cftimeindex_dayofyear_accessor(index): result = index.dayofyear expected = [date.dayofyr for date in index] assert_array_equal(result, expected) @requires_cftime def test_cftimeindex_dayofweek_accessor(index): result = index.dayofweek expected = [date.dayofwk for date in index] assert_array_equal(result, expected) @requires_cftime def test_cftimeindex_days_in_month_accessor(index): result = index.days_in_month expected = [date.daysinmonth for date in index] assert_array_equal(result, expected) @requires_cftime @pytest.mark.parametrize( ("string", "date_args", "reso"), [ ("1999", (1999, 1, 1), "year"), ("199902", (1999, 2, 1), "month"), ("19990202", (1999, 2, 2), "day"), ("19990202T01", (1999, 2, 2, 1), "hour"), ("19990202T0101", (1999, 2, 2, 1, 1), "minute"), ("19990202T010156", (1999, 2, 2, 1, 1, 56), "second"), ], ) def test_parse_iso8601_with_reso(date_type, string, date_args, reso): expected_date = date_type(*date_args) expected_reso = reso result_date, result_reso = _parse_iso8601_with_reso(date_type, string) assert result_date == expected_date assert result_reso == expected_reso @requires_cftime def test_parse_string_to_bounds_year(date_type, dec_days): parsed = date_type(2, 2, 10, 6, 2, 8, 1) expected_start = date_type(2, 1, 1) expected_end = date_type(2, 12, dec_days, 23, 59, 59, 999999) result_start, result_end = _parsed_string_to_bounds(date_type, "year", parsed) assert result_start == expected_start assert result_end == expected_end @requires_cftime def test_parse_string_to_bounds_month_feb(date_type, feb_days): parsed = date_type(2, 2, 10, 6, 2, 8, 1) expected_start = date_type(2, 2, 1) expected_end = date_type(2, 2, feb_days, 23, 59, 59, 999999) result_start, result_end = _parsed_string_to_bounds(date_type, "month", parsed) assert result_start == expected_start assert result_end == expected_end @requires_cftime def test_parse_string_to_bounds_month_dec(date_type, dec_days): parsed = date_type(2, 12, 1) expected_start = date_type(2, 12, 1) expected_end = date_type(2, 12, dec_days, 23, 59, 59, 999999) result_start, result_end = _parsed_string_to_bounds(date_type, "month", parsed) assert result_start == expected_start assert result_end == expected_end @requires_cftime @pytest.mark.parametrize( ("reso", "ex_start_args", "ex_end_args"), [ ("day", (2, 2, 10), (2, 2, 10, 23, 59, 59, 999999)), ("hour", (2, 2, 10, 6), (2, 2, 10, 6, 59, 59, 999999)), ("minute", (2, 2, 10, 6, 2), (2, 2, 10, 6, 2, 59, 999999)), ("second", (2, 2, 10, 6, 2, 8), (2, 2, 10, 6, 2, 8, 999999)), ], ) def test_parsed_string_to_bounds_sub_monthly( date_type, reso, ex_start_args, ex_end_args ): parsed = date_type(2, 2, 10, 6, 2, 8, 123456) expected_start = date_type(*ex_start_args) expected_end = date_type(*ex_end_args) result_start, result_end = _parsed_string_to_bounds(date_type, reso, parsed) assert result_start == expected_start assert result_end == expected_end @requires_cftime def test_parsed_string_to_bounds_raises(date_type): with pytest.raises(KeyError): _parsed_string_to_bounds(date_type, "a", date_type(1, 1, 1)) @requires_cftime def test_get_loc(date_type, index): result = index.get_loc("0001") assert result == slice(0, 2) result = index.get_loc(date_type(1, 2, 1)) assert result == 1 result = index.get_loc("0001-02-01") assert result == slice(1, 2) with pytest.raises(KeyError, match=r"1234"): index.get_loc("1234") @requires_cftime def test_get_slice_bound(date_type, index): # The kind argument is required in earlier versions of pandas even though it # is not used by CFTimeIndex. This logic can be removed once our minimum # version of pandas is at least 1.3. if Version(pd.__version__) < Version("1.3"): kind_args = ("getitem",) else: kind_args = () result = index.get_slice_bound("0001", "left", *kind_args) expected = 0 assert result == expected result = index.get_slice_bound("0001", "right", *kind_args) expected = 2 assert result == expected result = index.get_slice_bound(date_type(1, 3, 1), "left", *kind_args) expected = 2 assert result == expected result = index.get_slice_bound(date_type(1, 3, 1), "right", *kind_args) expected = 2 assert result == expected @requires_cftime def test_get_slice_bound_decreasing_index(date_type, monotonic_decreasing_index): # The kind argument is required in earlier versions of pandas even though it # is not used by CFTimeIndex. This logic can be removed once our minimum # version of pandas is at least 1.3. if Version(pd.__version__) < Version("1.3"): kind_args = ("getitem",) else: kind_args = () result = monotonic_decreasing_index.get_slice_bound("0001", "left", *kind_args) expected = 2 assert result == expected result = monotonic_decreasing_index.get_slice_bound("0001", "right", *kind_args) expected = 4 assert result == expected result = monotonic_decreasing_index.get_slice_bound( date_type(1, 3, 1), "left", *kind_args ) expected = 2 assert result == expected result = monotonic_decreasing_index.get_slice_bound( date_type(1, 3, 1), "right", *kind_args ) expected = 2 assert result == expected @requires_cftime def test_get_slice_bound_length_one_index(date_type, length_one_index): # The kind argument is required in earlier versions of pandas even though it # is not used by CFTimeIndex. This logic can be removed once our minimum # version of pandas is at least 1.3. if Version(pd.__version__) <= Version("1.3"): kind_args = ("getitem",) else: kind_args = () result = length_one_index.get_slice_bound("0001", "left", *kind_args) expected = 0 assert result == expected result = length_one_index.get_slice_bound("0001", "right", *kind_args) expected = 1 assert result == expected result = length_one_index.get_slice_bound(date_type(1, 3, 1), "left", *kind_args) expected = 1 assert result == expected result = length_one_index.get_slice_bound(date_type(1, 3, 1), "right", *kind_args) expected = 1 assert result == expected @requires_cftime def test_string_slice_length_one_index(length_one_index): da = xr.DataArray([1], coords=[length_one_index], dims=["time"]) result = da.sel(time=slice("0001", "0001")) assert_identical(result, da) @requires_cftime def test_date_type_property(date_type, index): assert index.date_type is date_type @requires_cftime def test_contains(date_type, index): assert "0001-01-01" in index assert "0001" in index assert "0003" not in index assert date_type(1, 1, 1) in index assert date_type(3, 1, 1) not in index @requires_cftime def test_groupby(da): result = da.groupby("time.month").sum("time") expected = xr.DataArray([4, 6], coords=[[1, 2]], dims=["month"]) assert_identical(result, expected) SEL_STRING_OR_LIST_TESTS = { "string": "0001", "string-slice": slice("0001-01-01", "0001-12-30"), "bool-list": [True, True, False, False], } @requires_cftime @pytest.mark.parametrize( "sel_arg", list(SEL_STRING_OR_LIST_TESTS.values()), ids=list(SEL_STRING_OR_LIST_TESTS.keys()), ) def test_sel_string_or_list(da, index, sel_arg): expected = xr.DataArray([1, 2], coords=[index[:2]], dims=["time"]) result = da.sel(time=sel_arg) assert_identical(result, expected) @requires_cftime def test_sel_date_slice_or_list(da, index, date_type): expected = xr.DataArray([1, 2], coords=[index[:2]], dims=["time"]) result = da.sel(time=slice(date_type(1, 1, 1), date_type(1, 12, 30))) assert_identical(result, expected) result = da.sel(time=[date_type(1, 1, 1), date_type(1, 2, 1)]) assert_identical(result, expected) @requires_cftime def test_sel_date_scalar(da, date_type, index): expected = xr.DataArray(1).assign_coords(time=index[0]) result = da.sel(time=date_type(1, 1, 1)) assert_identical(result, expected) @requires_cftime def test_sel_date_distant_date(da, date_type, index): expected = xr.DataArray(4).assign_coords(time=index[3]) result = da.sel(time=date_type(2000, 1, 1), method="nearest") assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [ {"method": "nearest"}, {"method": "nearest", "tolerance": timedelta(days=70)}, {"method": "nearest", "tolerance": timedelta(days=1800000)}, ], ) def test_sel_date_scalar_nearest(da, date_type, index, sel_kwargs): expected = xr.DataArray(2).assign_coords(time=index[1]) result = da.sel(time=date_type(1, 4, 1), **sel_kwargs) assert_identical(result, expected) expected = xr.DataArray(3).assign_coords(time=index[2]) result = da.sel(time=date_type(1, 11, 1), **sel_kwargs) assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [{"method": "pad"}, {"method": "pad", "tolerance": timedelta(days=365)}], ) def test_sel_date_scalar_pad(da, date_type, index, sel_kwargs): expected = xr.DataArray(2).assign_coords(time=index[1]) result = da.sel(time=date_type(1, 4, 1), **sel_kwargs) assert_identical(result, expected) expected = xr.DataArray(2).assign_coords(time=index[1]) result = da.sel(time=date_type(1, 11, 1), **sel_kwargs) assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [{"method": "backfill"}, {"method": "backfill", "tolerance": timedelta(days=365)}], ) def test_sel_date_scalar_backfill(da, date_type, index, sel_kwargs): expected = xr.DataArray(3).assign_coords(time=index[2]) result = da.sel(time=date_type(1, 4, 1), **sel_kwargs) assert_identical(result, expected) expected = xr.DataArray(3).assign_coords(time=index[2]) result = da.sel(time=date_type(1, 11, 1), **sel_kwargs) assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [ {"method": "pad", "tolerance": timedelta(days=20)}, {"method": "backfill", "tolerance": timedelta(days=20)}, {"method": "nearest", "tolerance": timedelta(days=20)}, ], ) def test_sel_date_scalar_tolerance_raises(da, date_type, sel_kwargs): with pytest.raises(KeyError): da.sel(time=date_type(1, 5, 1), **sel_kwargs) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [{"method": "nearest"}, {"method": "nearest", "tolerance": timedelta(days=70)}], ) def test_sel_date_list_nearest(da, date_type, index, sel_kwargs): expected = xr.DataArray([2, 2], coords=[[index[1], index[1]]], dims=["time"]) result = da.sel(time=[date_type(1, 3, 1), date_type(1, 4, 1)], **sel_kwargs) assert_identical(result, expected) expected = xr.DataArray([2, 3], coords=[[index[1], index[2]]], dims=["time"]) result = da.sel(time=[date_type(1, 3, 1), date_type(1, 12, 1)], **sel_kwargs) assert_identical(result, expected) expected = xr.DataArray([3, 3], coords=[[index[2], index[2]]], dims=["time"]) result = da.sel(time=[date_type(1, 11, 1), date_type(1, 12, 1)], **sel_kwargs) assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [{"method": "pad"}, {"method": "pad", "tolerance": timedelta(days=365)}], ) def test_sel_date_list_pad(da, date_type, index, sel_kwargs): expected = xr.DataArray([2, 2], coords=[[index[1], index[1]]], dims=["time"]) result = da.sel(time=[date_type(1, 3, 1), date_type(1, 4, 1)], **sel_kwargs) assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [{"method": "backfill"}, {"method": "backfill", "tolerance": timedelta(days=365)}], ) def test_sel_date_list_backfill(da, date_type, index, sel_kwargs): expected = xr.DataArray([3, 3], coords=[[index[2], index[2]]], dims=["time"]) result = da.sel(time=[date_type(1, 3, 1), date_type(1, 4, 1)], **sel_kwargs) assert_identical(result, expected) @requires_cftime @pytest.mark.parametrize( "sel_kwargs", [ {"method": "pad", "tolerance": timedelta(days=20)}, {"method": "backfill", "tolerance": timedelta(days=20)}, {"method": "nearest", "tolerance": timedelta(days=20)}, ], ) def test_sel_date_list_tolerance_raises(da, date_type, sel_kwargs): with pytest.raises(KeyError): da.sel(time=[date_type(1, 2, 1), date_type(1, 5, 1)], **sel_kwargs) @requires_cftime def test_isel(da, index): expected = xr.DataArray(1).assign_coords(time=index[0]) result = da.isel(time=0) assert_identical(result, expected) expected = xr.DataArray([1, 2], coords=[index[:2]], dims=["time"]) result = da.isel(time=[0, 1]) assert_identical(result, expected) @pytest.fixture def scalar_args(date_type): return [date_type(1, 1, 1)] @pytest.fixture def range_args(date_type): return [ "0001", slice("0001-01-01", "0001-12-30"), slice(None, "0001-12-30"), slice(date_type(1, 1, 1), date_type(1, 12, 30)), slice(None, date_type(1, 12, 30)), ] @requires_cftime def test_indexing_in_series_getitem(series, index, scalar_args, range_args): for arg in scalar_args: assert series[arg] == 1 expected = pd.Series([1, 2], index=index[:2]) for arg in range_args: assert series[arg].equals(expected) @requires_cftime def test_indexing_in_series_loc(series, index, scalar_args, range_args): for arg in scalar_args: assert series.loc[arg] == 1 expected = pd.Series([1, 2], index=index[:2]) for arg in range_args: assert series.loc[arg].equals(expected) @requires_cftime def test_indexing_in_series_iloc(series, index): expected = 1 assert series.iloc[0] == expected expected = pd.Series([1, 2], index=index[:2]) assert series.iloc[:2].equals(expected) @requires_cftime def test_series_dropna(index): series = pd.Series([0.0, 1.0, np.nan, np.nan], index=index) expected = series.iloc[:2] result = series.dropna() assert result.equals(expected) @requires_cftime def test_indexing_in_dataframe_loc(df, index, scalar_args, range_args): expected = pd.Series([1], name=index[0]) for arg in scalar_args: result = df.loc[arg] assert result.equals(expected) expected = pd.DataFrame([1, 2], index=index[:2]) for arg in range_args: result = df.loc[arg] assert result.equals(expected) @requires_cftime def test_indexing_in_dataframe_iloc(df, index): expected = pd.Series([1], name=index[0]) result = df.iloc[0] assert result.equals(expected) assert result.equals(expected) expected = pd.DataFrame([1, 2], index=index[:2]) result = df.iloc[:2] assert result.equals(expected) @requires_cftime def test_concat_cftimeindex(date_type): da1 = xr.DataArray( [1.0, 2.0], coords=[[date_type(1, 1, 1), date_type(1, 2, 1)]], dims=["time"] ) da2 = xr.DataArray( [3.0, 4.0], coords=[[date_type(1, 3, 1), date_type(1, 4, 1)]], dims=["time"] ) da = xr.concat([da1, da2], dim="time") assert isinstance(da.xindexes["time"].to_pandas_index(), CFTimeIndex) @requires_cftime def test_empty_cftimeindex(): index = CFTimeIndex([]) assert index.date_type is None @requires_cftime def test_cftimeindex_add(index): date_type = index.date_type expected_dates = [ date_type(1, 1, 2), date_type(1, 2, 2), date_type(2, 1, 2), date_type(2, 2, 2), ] expected = CFTimeIndex(expected_dates) result = index + timedelta(days=1) assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_add_timedeltaindex(calendar) -> None: a = xr.cftime_range("2000", periods=5, calendar=calendar) deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) result = a + deltas expected = a.shift(2, "D") assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime @pytest.mark.parametrize("n", [2.0, 1.5]) @pytest.mark.parametrize( "freq,units", [ ("D", "D"), ("H", "H"), ("T", "min"), ("S", "S"), ("L", "ms"), ], ) @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_shift_float(n, freq, units, calendar) -> None: a = xr.cftime_range("2000", periods=3, calendar=calendar, freq="D") result = a + pd.Timedelta(n, units) expected = a.shift(n, freq) assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime def test_cftimeindex_shift_float_us() -> None: a = xr.cftime_range("2000", periods=3, freq="D") with pytest.raises( ValueError, match="Could not convert to integer offset at any resolution" ): a.shift(2.5, "us") @requires_cftime @pytest.mark.parametrize("freq", ["AS", "A", "YS", "Y", "QS", "Q", "MS", "M"]) def test_cftimeindex_shift_float_fails_for_non_tick_freqs(freq) -> None: a = xr.cftime_range("2000", periods=3, freq="D") with pytest.raises(TypeError, match="unsupported operand type"): a.shift(2.5, freq) @requires_cftime def test_cftimeindex_radd(index): date_type = index.date_type expected_dates = [ date_type(1, 1, 2), date_type(1, 2, 2), date_type(2, 1, 2), date_type(2, 2, 2), ] expected = CFTimeIndex(expected_dates) result = timedelta(days=1) + index assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_timedeltaindex_add_cftimeindex(calendar) -> None: a = xr.cftime_range("2000", periods=5, calendar=calendar) deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) result = deltas + a expected = a.shift(2, "D") assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime def test_cftimeindex_sub_timedelta(index): date_type = index.date_type expected_dates = [ date_type(1, 1, 2), date_type(1, 2, 2), date_type(2, 1, 2), date_type(2, 2, 2), ] expected = CFTimeIndex(expected_dates) result = index + timedelta(days=2) result = result - timedelta(days=1) assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime @pytest.mark.parametrize( "other", [np.array(4 * [timedelta(days=1)]), np.array(timedelta(days=1))], ids=["1d-array", "scalar-array"], ) def test_cftimeindex_sub_timedelta_array(index, other): date_type = index.date_type expected_dates = [ date_type(1, 1, 2), date_type(1, 2, 2), date_type(2, 1, 2), date_type(2, 2, 2), ] expected = CFTimeIndex(expected_dates) result = index + timedelta(days=2) result = result - other assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_sub_cftimeindex(calendar) -> None: a = xr.cftime_range("2000", periods=5, calendar=calendar) b = a.shift(2, "D") result = b - a expected = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) assert result.equals(expected) assert isinstance(result, pd.TimedeltaIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_sub_cftime_datetime(calendar): a = xr.cftime_range("2000", periods=5, calendar=calendar) result = a - a[0] expected = pd.TimedeltaIndex([timedelta(days=i) for i in range(5)]) assert result.equals(expected) assert isinstance(result, pd.TimedeltaIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftime_datetime_sub_cftimeindex(calendar): a = xr.cftime_range("2000", periods=5, calendar=calendar) result = a[0] - a expected = pd.TimedeltaIndex([timedelta(days=-i) for i in range(5)]) assert result.equals(expected) assert isinstance(result, pd.TimedeltaIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_distant_cftime_datetime_sub_cftimeindex(calendar): a = xr.cftime_range("2000", periods=5, calendar=calendar) with pytest.raises(ValueError, match="difference exceeds"): a.date_type(1, 1, 1) - a @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_sub_timedeltaindex(calendar) -> None: a = xr.cftime_range("2000", periods=5, calendar=calendar) deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) result = a - deltas expected = a.shift(-2, "D") assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_sub_index_of_cftime_datetimes(calendar): a = xr.cftime_range("2000", periods=5, calendar=calendar) b = pd.Index(a.values) expected = a - a result = a - b assert result.equals(expected) assert isinstance(result, pd.TimedeltaIndex) @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_cftimeindex_sub_not_implemented(calendar): a = xr.cftime_range("2000", periods=5, calendar=calendar) with pytest.raises(TypeError, match="unsupported operand"): a - 1 @requires_cftime def test_cftimeindex_rsub(index): with pytest.raises(TypeError): timedelta(days=1) - index @requires_cftime @pytest.mark.parametrize("freq", ["D", timedelta(days=1)]) def test_cftimeindex_shift(index, freq) -> None: date_type = index.date_type expected_dates = [ date_type(1, 1, 3), date_type(1, 2, 3), date_type(2, 1, 3), date_type(2, 2, 3), ] expected = CFTimeIndex(expected_dates) result = index.shift(2, freq) assert result.equals(expected) assert isinstance(result, CFTimeIndex) @requires_cftime def test_cftimeindex_shift_invalid_n() -> None: index = xr.cftime_range("2000", periods=3) with pytest.raises(TypeError): index.shift("a", "D") @requires_cftime def test_cftimeindex_shift_invalid_freq() -> None: index = xr.cftime_range("2000", periods=3) with pytest.raises(TypeError): index.shift(1, 1) @requires_cftime @pytest.mark.parametrize( ("calendar", "expected"), [ ("noleap", "noleap"), ("365_day", "noleap"), ("360_day", "360_day"), ("julian", "julian"), ("gregorian", standard_or_gregorian), ("standard", standard_or_gregorian), ("proleptic_gregorian", "proleptic_gregorian"), ], ) def test_cftimeindex_calendar_property(calendar, expected): index = xr.cftime_range(start="2000", periods=3, calendar=calendar) assert index.calendar == expected @requires_cftime @pytest.mark.parametrize( ("calendar", "expected"), [ ("noleap", "noleap"), ("365_day", "noleap"), ("360_day", "360_day"), ("julian", "julian"), ("gregorian", standard_or_gregorian), ("standard", standard_or_gregorian), ("proleptic_gregorian", "proleptic_gregorian"), ], ) def test_cftimeindex_calendar_repr(calendar, expected): """Test that cftimeindex has calendar property in repr.""" index = xr.cftime_range(start="2000", periods=3, calendar=calendar) repr_str = index.__repr__() assert f" calendar='{expected}'" in repr_str assert "2000-01-01 00:00:00, 2000-01-02 00:00:00" in repr_str @requires_cftime @pytest.mark.parametrize("periods", [2, 40]) def test_cftimeindex_periods_repr(periods): """Test that cftimeindex has periods property in repr.""" index = xr.cftime_range(start="2000", periods=periods) repr_str = index.__repr__() assert f" length={periods}" in repr_str @requires_cftime @pytest.mark.parametrize("calendar", ["noleap", "360_day", "standard"]) @pytest.mark.parametrize("freq", ["D", "H"]) def test_cftimeindex_freq_in_repr(freq, calendar): """Test that cftimeindex has frequency property in repr.""" index = xr.cftime_range(start="2000", periods=3, freq=freq, calendar=calendar) repr_str = index.__repr__() assert f", freq='{freq}'" in repr_str @requires_cftime @pytest.mark.parametrize( "periods,expected", [ ( 2, f"""\ CFTimeIndex([2000-01-01 00:00:00, 2000-01-02 00:00:00], dtype='object', length=2, calendar='{standard_or_gregorian}', freq=None)""", ), ( 4, f"""\ CFTimeIndex([2000-01-01 00:00:00, 2000-01-02 00:00:00, 2000-01-03 00:00:00, 2000-01-04 00:00:00], dtype='object', length=4, calendar='{standard_or_gregorian}', freq='D')""", ), ( 101, f"""\ CFTimeIndex([2000-01-01 00:00:00, 2000-01-02 00:00:00, 2000-01-03 00:00:00, 2000-01-04 00:00:00, 2000-01-05 00:00:00, 2000-01-06 00:00:00, 2000-01-07 00:00:00, 2000-01-08 00:00:00, 2000-01-09 00:00:00, 2000-01-10 00:00:00, ... 2000-04-01 00:00:00, 2000-04-02 00:00:00, 2000-04-03 00:00:00, 2000-04-04 00:00:00, 2000-04-05 00:00:00, 2000-04-06 00:00:00, 2000-04-07 00:00:00, 2000-04-08 00:00:00, 2000-04-09 00:00:00, 2000-04-10 00:00:00], dtype='object', length=101, calendar='{standard_or_gregorian}', freq='D')""", ), ], ) def test_cftimeindex_repr_formatting(periods, expected): """Test that cftimeindex.__repr__ is formatted similar to pd.Index.__repr__.""" index = xr.cftime_range(start="2000", periods=periods, freq="D") expected = dedent(expected) assert expected == repr(index) @requires_cftime @pytest.mark.parametrize("display_width", [40, 80, 100]) @pytest.mark.parametrize("periods", [2, 3, 4, 100, 101]) def test_cftimeindex_repr_formatting_width(periods, display_width): """Test that cftimeindex is sensitive to OPTIONS['display_width'].""" index = xr.cftime_range(start="2000", periods=periods) len_intro_str = len("CFTimeIndex(") with xr.set_options(display_width=display_width): repr_str = index.__repr__() splitted = repr_str.split("\n") for i, s in enumerate(splitted): # check that lines not longer than OPTIONS['display_width'] assert len(s) <= display_width, f"{len(s)} {s} {display_width}" if i > 0: # check for initial spaces assert s[:len_intro_str] == " " * len_intro_str @requires_cftime @pytest.mark.parametrize("periods", [22, 50, 100]) def test_cftimeindex_repr_101_shorter(periods): index_101 = xr.cftime_range(start="2000", periods=101) index_periods = xr.cftime_range(start="2000", periods=periods) index_101_repr_str = index_101.__repr__() index_periods_repr_str = index_periods.__repr__() assert len(index_101_repr_str) < len(index_periods_repr_str) @requires_cftime def test_parse_array_of_cftime_strings(): from cftime import DatetimeNoLeap strings = np.array([["2000-01-01", "2000-01-02"], ["2000-01-03", "2000-01-04"]]) expected = np.array( [ [DatetimeNoLeap(2000, 1, 1), DatetimeNoLeap(2000, 1, 2)], [DatetimeNoLeap(2000, 1, 3), DatetimeNoLeap(2000, 1, 4)], ] ) result = _parse_array_of_cftime_strings(strings, DatetimeNoLeap) np.testing.assert_array_equal(result, expected) # Test scalar array case strings = np.array("2000-01-01") expected = np.array(DatetimeNoLeap(2000, 1, 1)) result = _parse_array_of_cftime_strings(strings, DatetimeNoLeap) np.testing.assert_array_equal(result, expected) @requires_cftime @pytest.mark.parametrize("calendar", _ALL_CALENDARS) def test_strftime_of_cftime_array(calendar): date_format = "%Y%m%d%H%M" cf_values = xr.cftime_range("2000", periods=5, calendar=calendar) dt_values = pd.date_range("2000", periods=5) expected = pd.Index(dt_values.strftime(date_format)) result = cf_values.strftime(date_format) assert result.equals(expected) @requires_cftime @pytest.mark.parametrize("calendar", _ALL_CALENDARS) @pytest.mark.parametrize("unsafe", [False, True]) def test_to_datetimeindex(calendar, unsafe): index = xr.cftime_range("2000", periods=5, calendar=calendar) expected = pd.date_range("2000", periods=5) if calendar in _NON_STANDARD_CALENDARS and not unsafe: with pytest.warns(RuntimeWarning, match="non-standard"): result = index.to_datetimeindex() else: result = index.to_datetimeindex(unsafe=unsafe) assert result.equals(expected) np.testing.assert_array_equal(result, expected) assert isinstance(result, pd.DatetimeIndex) @requires_cftime @pytest.mark.parametrize("calendar", _ALL_CALENDARS) def test_to_datetimeindex_out_of_range(calendar): index = xr.cftime_range("0001", periods=5, calendar=calendar) with pytest.raises(ValueError, match="0001"): index.to_datetimeindex() @requires_cftime @pytest.mark.parametrize("calendar", ["all_leap", "360_day"]) def test_to_datetimeindex_feb_29(calendar): index = xr.cftime_range("2001-02-28", periods=2, calendar=calendar) with pytest.raises(ValueError, match="29"): index.to_datetimeindex() @requires_cftime @pytest.mark.xfail(reason="https://github.com/pandas-dev/pandas/issues/24263") def test_multiindex(): index = xr.cftime_range("2001-01-01", periods=100, calendar="360_day") mindex = pd.MultiIndex.from_arrays([index]) assert mindex.get_loc("2001-01") == slice(0, 30) @requires_cftime @pytest.mark.parametrize("freq", ["3663S", "33T", "2H"]) @pytest.mark.parametrize("method", ["floor", "ceil", "round"]) def test_rounding_methods_against_datetimeindex(freq, method): expected = pd.date_range("2000-01-02T01:03:51", periods=10, freq="1777S") expected = getattr(expected, method)(freq) result = xr.cftime_range("2000-01-02T01:03:51", periods=10, freq="1777S") result = getattr(result, method)(freq).to_datetimeindex() assert result.equals(expected) @requires_cftime @pytest.mark.parametrize("method", ["floor", "ceil", "round"]) def test_rounding_methods_invalid_freq(method): index = xr.cftime_range("2000-01-02T01:03:51", periods=10, freq="1777S") with pytest.raises(ValueError, match="fixed"): getattr(index, method)("MS") @pytest.fixture def rounding_index(date_type): return xr.CFTimeIndex( [ date_type(1, 1, 1, 1, 59, 59, 999512), date_type(1, 1, 1, 3, 0, 1, 500001), date_type(1, 1, 1, 7, 0, 6, 499999), ] ) @requires_cftime def test_ceil(rounding_index, date_type): result = rounding_index.ceil("S") expected = xr.CFTimeIndex( [ date_type(1, 1, 1, 2, 0, 0, 0), date_type(1, 1, 1, 3, 0, 2, 0), date_type(1, 1, 1, 7, 0, 7, 0), ] ) assert result.equals(expected) @requires_cftime def test_floor(rounding_index, date_type): result = rounding_index.floor("S") expected = xr.CFTimeIndex( [ date_type(1, 1, 1, 1, 59, 59, 0), date_type(1, 1, 1, 3, 0, 1, 0), date_type(1, 1, 1, 7, 0, 6, 0), ] ) assert result.equals(expected) @requires_cftime def test_round(rounding_index, date_type): result = rounding_index.round("S") expected = xr.CFTimeIndex( [ date_type(1, 1, 1, 2, 0, 0, 0), date_type(1, 1, 1, 3, 0, 2, 0), date_type(1, 1, 1, 7, 0, 6, 0), ] ) assert result.equals(expected) @requires_cftime def test_asi8(date_type): index = xr.CFTimeIndex([date_type(1970, 1, 1), date_type(1970, 1, 2)]) result = index.asi8 expected = 1000000 * 86400 * np.array([0, 1]) np.testing.assert_array_equal(result, expected) @requires_cftime def test_asi8_distant_date(): """Test that asi8 conversion is truly exact.""" import cftime date_type = cftime.DatetimeProlepticGregorian index = xr.CFTimeIndex([date_type(10731, 4, 22, 3, 25, 45, 123456)]) result = index.asi8 expected = np.array([1000000 * 86400 * 400 * 8000 + 12345 * 1000000 + 123456]) np.testing.assert_array_equal(result, expected) @requires_cftime def test_infer_freq_valid_types(): cf_indx = xr.cftime_range("2000-01-01", periods=3, freq="D") assert xr.infer_freq(cf_indx) == "D" assert xr.infer_freq(xr.DataArray(cf_indx)) == "D" pd_indx = pd.date_range("2000-01-01", periods=3, freq="D") assert xr.infer_freq(pd_indx) == "D" assert xr.infer_freq(xr.DataArray(pd_indx)) == "D" pd_td_indx = pd.timedelta_range(start="1D", periods=3, freq="D") assert xr.infer_freq(pd_td_indx) == "D" assert xr.infer_freq(xr.DataArray(pd_td_indx)) == "D" @requires_cftime def test_infer_freq_invalid_inputs(): # Non-datetime DataArray with pytest.raises(ValueError, match="must contain datetime-like objects"): xr.infer_freq(xr.DataArray([0, 1, 2])) indx = xr.cftime_range("1990-02-03", periods=4, freq="MS") # 2D DataArray with pytest.raises(ValueError, match="must be 1D"): xr.infer_freq(xr.DataArray([indx, indx])) # CFTimeIndex too short with pytest.raises(ValueError, match="Need at least 3 dates to infer frequency"): xr.infer_freq(indx[:2]) # Non-monotonic input assert xr.infer_freq(indx[np.array([0, 2, 1, 3])]) is None # Non-unique input assert xr.infer_freq(indx[np.array([0, 1, 1, 2])]) is None # No unique frequency (here 1st step is MS, second is 2MS) assert xr.infer_freq(indx[np.array([0, 1, 3])]) is None # Same, but for QS indx = xr.cftime_range("1990-02-03", periods=4, freq="QS") assert xr.infer_freq(indx[np.array([0, 1, 3])]) is None @requires_cftime @pytest.mark.parametrize( "freq", [ "300AS-JAN", "A-DEC", "AS-JUL", "2AS-FEB", "Q-NOV", "3QS-DEC", "MS", "4M", "7D", "D", "30H", "5T", "40S", ], ) @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_infer_freq(freq, calendar): indx = xr.cftime_range("2000-01-01", periods=3, freq=freq, calendar=calendar) out = xr.infer_freq(indx) assert out == freq @requires_cftime @pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) def test_pickle_cftimeindex(calendar): idx = xr.cftime_range("2000-01-01", periods=3, freq="D", calendar=calendar) idx_pkl = pickle.loads(pickle.dumps(idx)) assert (idx == idx_pkl).all()