WIP: Support calendar-specific cftime.datetime instances

xarray/coding/cftime_offsets.py

@@ -77,7 +77,7 @@
     from xarray.core.types import InclusiveOptions, SideOptions
 
 
-def get_date_type(calendar, use_cftime=True):
+def get_date_type(calendar, use_cftime=True, legacy=True):
     """Return the cftime date type for a given calendar name."""
     if cftime is None:
         raise ImportError("cftime is required for dates with non-standard calendars")
@@ -96,7 +96,10 @@ def get_date_type(calendar, use_cftime=True):
             "all_leap": cftime.DatetimeAllLeap,
             "standard": cftime.DatetimeGregorian,
         }
-        return calendars[calendar]
+        if legacy:
+            return calendars[calendar]
+        else:
+            return partial(cftime.datetime, calendar=calendar)
 
 
 class BaseCFTimeOffset:

xarray/coding/cftimeindex.py

@@ -45,6 +45,7 @@
 import re
 import warnings
 from datetime import timedelta
+from functools import partial
 
 import numpy as np
 import pandas as pd
@@ -213,8 +214,15 @@ def f(self, min_cftime_version=min_cftime_version):
 
 
 def get_date_type(self):
+    if cftime is None:
+        raise ModuleNotFoundError("No module named 'cftime'")
+
     if self._data.size:
-        return type(self._data[0])
+        sample = self._data[0]
+        if type(sample) is cftime.datetime:
+            return partial(cftime.datetime, calendar=sample.calendar)
+        else:
+            return type(sample)
     else:
         return None
 
@@ -231,10 +239,13 @@ def assert_all_valid_date_type(data):
                 "CFTimeIndex requires cftime.datetime "
                 f"objects. Got object of {date_type}."
             )
-        if not all(isinstance(value, date_type) for value in data):
+        calendar = sample.calendar
+        if not all(
+            hasattr(value, "calendar") and value.calendar == calendar for value in data
+        ):
             raise TypeError(
                 "CFTimeIndex requires using datetime "
-                f"objects of all the same type.  Got\n{data}."
+                f"objects with the same calendar. Got\n{data}."
             )
 
 
@@ -682,17 +693,12 @@ def strftime(self, date_format):
     @property
     def asi8(self):
         """Convert to integers with units of microseconds since 1970-01-01."""
-        from xarray.core.resample_cftime import exact_cftime_datetime_difference
-
         if not self._data.size:
             return np.array([], dtype=np.int64)
 
         epoch = self.date_type(1970, 1, 1)
         return np.array(
-            [
-                _total_microseconds(exact_cftime_datetime_difference(epoch, date))
-                for date in self.values
-            ],
+            [_total_microseconds(date - epoch) for date in self.values],
             dtype=np.int64,
         )
 

xarray/core/resample_cftime.py

@@ -454,46 +454,6 @@ def _adjust_dates_anchored(
     return fresult, lresult
 
 
-def exact_cftime_datetime_difference(a: CFTimeDatetime, b: CFTimeDatetime):
-    """Exact computation of b - a
-
-    Assumes:
-
-        a = a_0 + a_m
-        b = b_0 + b_m
-
-    Here a_0, and b_0 represent the input dates rounded
-    down to the nearest second, and a_m, and b_m represent
-    the remaining microseconds associated with date a and
-    date b.
-
-    We can then express the value of b - a as:
-
-        b - a = (b_0 + b_m) - (a_0 + a_m) = b_0 - a_0 + b_m - a_m
-
-    By construction, we know that b_0 - a_0 must be a round number
-    of seconds.  Therefore we can take the result of b_0 - a_0 using
-    ordinary cftime.datetime arithmetic and round to the nearest
-    second.  b_m - a_m is the remainder, in microseconds, and we
-    can simply add this to the rounded timedelta.
-
-    Parameters
-    ----------
-    a : cftime.datetime
-        Input datetime
-    b : cftime.datetime
-        Input datetime
-
-    Returns
-    -------
-    datetime.timedelta
-    """
-    seconds = b.replace(microsecond=0) - a.replace(microsecond=0)
-    seconds = int(round(seconds.total_seconds()))
-    microseconds = b.microsecond - a.microsecond
-    return datetime.timedelta(seconds=seconds, microseconds=microseconds)
-
-
 def _convert_offset_to_timedelta(
     offset: datetime.timedelta | str | BaseCFTimeOffset,
 ) -> datetime.timedelta:

xarray/tests/test_cftime_offsets.py

@@ -1244,10 +1244,19 @@ def test_rollback(calendar, offset, initial_date_args, partial_expected_date_arg
     _CFTIME_RANGE_TESTS,
     ids=_id_func,
 )
+@pytest.mark.parametrize("use_legacy_date_type", [True, False])
 def test_cftime_range(
-    start, end, periods, freq, inclusive, normalize, calendar, expected_date_args
+    start,
+    end,
+    periods,
+    freq,
+    inclusive,
+    normalize,
+    calendar,
+    expected_date_args,
+    use_legacy_date_type,
 ):
-    date_type = get_date_type(calendar)
+    date_type = get_date_type(calendar, legacy=use_legacy_date_type)
     expected_dates = [date_type(*args) for args in expected_date_args]
 
     if isinstance(start, tuple):
@@ -1267,17 +1276,7 @@ def test_cftime_range(
     resulting_dates = result.values
 
     assert isinstance(result, CFTimeIndex)
-
-    if freq is not None:
-        np.testing.assert_equal(resulting_dates, expected_dates)
-    else:
-        # If we create a linear range of dates using cftime.num2date
-        # we will not get exact round number dates.  This is because
-        # datetime arithmetic in cftime is accurate approximately to
-        # 1 millisecond (see https://unidata.github.io/cftime/api.html).
-        deltas = resulting_dates - expected_dates
-        deltas = np.array([delta.total_seconds() for delta in deltas])
-        assert np.max(np.abs(deltas)) < 0.001
+    np.testing.assert_equal(resulting_dates, expected_dates)
 
 
 def test_cftime_range_name():

xarray/tests/test_cftimeindex.py

@@ -1306,10 +1306,18 @@ def test_asi8_empty_cftimeindex():
 
 @requires_cftime
 def test_infer_freq_valid_types():
+    import cftime
+
     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"
 
+    cf_indx = xr.cftime_range(
+        cftime.datetime(2000, 1, 1, calendar="noleap"), 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"

xarray/tests/test_coding_times.py

@@ -2,6 +2,7 @@
 
 import warnings
 from datetime import timedelta
+from functools import partial
 from itertools import product
 
 import numpy as np
@@ -102,7 +103,7 @@ def _all_cftime_date_types():
     import cftime
 
     return {
-        "noleap": cftime.DatetimeNoLeap,
+        "noleap": partial(cftime.datetime, calendar="noleap"),
         "365_day": cftime.DatetimeNoLeap,
         "360_day": cftime.Datetime360Day,
         "julian": cftime.DatetimeJulian,