From ff15a08bea27674923afa494b303c6e5cb4d513c Mon Sep 17 00:00:00 2001
From: Mark Harfouche <mark.harfouche@gmail.com>
Date: Wed, 10 Jul 2024 22:00:09 -0400
Subject: [PATCH] Fix time indexing regression in `convert_calendar` (#9192)

* MRC -- Selecting with string for cftime

See discussion in #9138

This commit and pull request mostly serves as a staging group for a
potential fix.

Test with:

```
pytest xarray/tests/test_cftimeindex.py::test_cftime_noleap_with_str
```

* effectively remove fastpath

* Add docstring

* Revert "effectively remove fastpath"

This reverts commit 0f1a5a2271e5522b5dd946d7f4f38f591211286e.

* Fix by reassigning coordinate

* Update what's new entry

* Simplify if condition

---------

Co-authored-by: Spencer Clark <spencerkclark@gmail.com>
---
 doc/whats-new.rst                 |  6 ++++++
 xarray/coding/calendar_ops.py     | 12 +++++++++++-
 xarray/tests/test_calendar_ops.py | 25 ++++++++++++++++++++++++-
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 0c401c2348e..f237b406bd5 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -47,6 +47,12 @@ Bug fixes
   By `Justus Magin <https://github.com/keewis>`_.
 - Promote floating-point numeric datetimes before decoding (:issue:`9179`, :pull:`9182`).
   By `Justus Magin <https://github.com/keewis>`_.
+- Address regression introduced in :pull:`9002` that prevented objects returned
+  by py:meth:`DataArray.convert_calendar` to be indexed by a time index in
+  certain circumstances (:issue:`9138`, :pull:`9192`). By `Mark Harfouche
+  <https://github.com/hmaarrfk>`_ and `Spencer Clark
+  <https://github.com/spencerkclark>`.
+
 - Fiy static typing of tolerance arguments by allowing `str` type (:issue:`8892`, :pull:`9194`).
   By `Michael Niklas <https://github.com/headtr1ck>`_.
 - Dark themes are now properly detected for ``html[data-theme=dark]``-tags (:pull:`9200`).
diff --git a/xarray/coding/calendar_ops.py b/xarray/coding/calendar_ops.py
index c4fe9e1f4ae..6f492e78bf9 100644
--- a/xarray/coding/calendar_ops.py
+++ b/xarray/coding/calendar_ops.py
@@ -5,7 +5,10 @@
 
 from xarray.coding.cftime_offsets import date_range_like, get_date_type
 from xarray.coding.cftimeindex import CFTimeIndex
-from xarray.coding.times import _should_cftime_be_used, convert_times
+from xarray.coding.times import (
+    _should_cftime_be_used,
+    convert_times,
+)
 from xarray.core.common import _contains_datetime_like_objects, is_np_datetime_like
 
 try:
@@ -222,6 +225,13 @@ def convert_calendar(
         # Remove NaN that where put on invalid dates in target calendar
         out = out.where(out[dim].notnull(), drop=True)
 
+        if use_cftime:
+            # Reassign times to ensure time index of output is a CFTimeIndex
+            # (previously it was an Index due to the presence of NaN values).
+            # Note this is not needed in the case that the output time index is
+            # a DatetimeIndex, since DatetimeIndexes can handle NaN values.
+            out[dim] = CFTimeIndex(out[dim].data)
+
     if missing is not None:
         time_target = date_range_like(time, calendar=calendar, use_cftime=use_cftime)
         out = out.reindex({dim: time_target}, fill_value=missing)
diff --git a/xarray/tests/test_calendar_ops.py b/xarray/tests/test_calendar_ops.py
index 7d229371808..13e9f7a1030 100644
--- a/xarray/tests/test_calendar_ops.py
+++ b/xarray/tests/test_calendar_ops.py
@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import numpy as np
+import pandas as pd
 import pytest
 
-from xarray import DataArray, infer_freq
+from xarray import CFTimeIndex, DataArray, infer_freq
 from xarray.coding.calendar_ops import convert_calendar, interp_calendar
 from xarray.coding.cftime_offsets import date_range
 from xarray.testing import assert_identical
@@ -286,3 +287,25 @@ def test_interp_calendar_errors():
         ValueError, match="Both 'source.x' and 'target' must contain datetime objects."
     ):
         interp_calendar(da1, da2, dim="x")
+
+
+@requires_cftime
+@pytest.mark.parametrize(
+    ("source_calendar", "target_calendar", "expected_index"),
+    [("standard", "noleap", CFTimeIndex), ("all_leap", "standard", pd.DatetimeIndex)],
+)
+def test_convert_calendar_produces_time_index(
+    source_calendar, target_calendar, expected_index
+):
+    # https://github.com/pydata/xarray/issues/9138
+    time = date_range("2000-01-01", "2002-01-01", freq="D", calendar=source_calendar)
+    temperature = np.ones(len(time))
+    da = DataArray(
+        data=temperature,
+        dims=["time"],
+        coords=dict(
+            time=time,
+        ),
+    )
+    converted = da.convert_calendar(target_calendar)
+    assert isinstance(converted.indexes["time"], expected_index)