_season_from_months can now handle np.nan (pydata#5876)

* _season_from_months can now handle np.nan

_season_from_months can now handle np.nan and values outside of [1,12]

I passed these tests:

def test_season():
    months = np.array([ 1, 2, 3, 4, 5, np.nan])
    assert ( _season_from_months(months) == np.array(['DJF', 'DJF', 'MAM', 'MAM', 'MAM', 'na']) ).all()
    
    months = np.array([ 1, 100, 3, 13, 0, -5])
    assert ( _season_from_months(months) == np.array(['DJF', 'na', 'MAM', 'na', 'na', 'na']) ).all()
    
    months = np.array(range(1, 13))
    assert ( _season_from_months(months) == np.array(['DJF', 'DJF', 'MAM', 'MAM', 'MAM', 'JJA', 'JJA', 'JJA', 'SON', 'SON', 'SON', 'DJF']) ).all()


test_season()

* Run black

* Remove duplicated import

* returns np.nan and removed the useless attribution

returns np.nan and removed the useless attribution when month is not in [1,12]

* added test

* applied black recommendations

* Apply suggestions from code review

finally, NaT is output as a "nan" string

Co-authored-by: Spencer Clark <spencerkclark@gmail.com>

* update whatsnew

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Correction on the credits

Co-authored-by: Spencer Clark <spencerkclark@gmail.com>

Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Spencer Clark <spencerkclark@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

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 <https://github.com/tcchiao>`_.
 
+- `dt.season <https://xarray.pydata.org/en/stable/generated/xarray.DataArray.dt.season.html>`_  can now handle NaN and NaT.  (:pull:`5876`).
+  By `Pierre Loicq <https://github.com/pierreloicq>`_.
+
+
 Documentation
 ~~~~~~~~~~~~~
 

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):

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",
             ]
         )