diff --git a/doc/whats-new.rst b/doc/whats-new.rst index f41e57989be..379c2c9d947 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -53,6 +53,10 @@ Bug fixes - Fix applying function with non-xarray arguments using :py:func:`xr.map_blocks`. By `Cindy Chiao `_. +- `dt.season `_ can now handle NaN and NaT. (:pull:`5876`). + By `Pierre Loicq `_. + + Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/accessor_dt.py b/xarray/core/accessor_dt.py index 2a7b6200d3b..7f8bf79a50a 100644 --- a/xarray/core/accessor_dt.py +++ b/xarray/core/accessor_dt.py @@ -16,9 +16,20 @@ def _season_from_months(months): """Compute season (DJF, MAM, JJA, SON) from month ordinal""" # TODO: Move "season" accessor upstream into pandas - seasons = np.array(["DJF", "MAM", "JJA", "SON"]) + seasons = np.array(["DJF", "MAM", "JJA", "SON", "nan"]) months = np.asarray(months) - return seasons[(months // 3) % 4] + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", message="invalid value encountered in floor_divide" + ) + warnings.filterwarnings( + "ignore", message="invalid value encountered in remainder" + ) + idx = (months // 3) % 4 + + idx[np.isnan(idx)] = 4 + return seasons[idx.astype(int)] def _access_through_cftimeindex(values, name): diff --git a/xarray/tests/test_accessor_dt.py b/xarray/tests/test_accessor_dt.py index b471bd2e267..e9278f1e918 100644 --- a/xarray/tests/test_accessor_dt.py +++ b/xarray/tests/test_accessor_dt.py @@ -222,6 +222,7 @@ def test_dask_accessor_method(self, method, parameters) -> None: def test_seasons(self) -> None: dates = pd.date_range(start="2000/01/01", freq="M", periods=12) + dates = dates.append(pd.Index([np.datetime64("NaT")])) dates = xr.DataArray(dates) seasons = xr.DataArray( [ @@ -237,6 +238,7 @@ def test_seasons(self) -> None: "SON", "SON", "DJF", + "nan", ] )