REF: tighten what we accept in TimedeltaIndex._simple_new (#31315)

Author: jbrockmendel

pandas/core/arrays/timedeltas.py

@@ -195,9 +195,12 @@ def __init__(self, values, dtype=_TD_DTYPE, freq=None, copy=False):
     def _simple_new(cls, values, freq=None, dtype=_TD_DTYPE):
         assert dtype == _TD_DTYPE, dtype
         assert isinstance(values, np.ndarray), type(values)
+        if values.dtype != _TD_DTYPE:
+            assert values.dtype == "i8"
+            values = values.view(_TD_DTYPE)
 
         result = object.__new__(cls)
-        result._data = values.view(_TD_DTYPE)
+        result._data = values
         result._freq = to_offset(freq)
         result._dtype = _TD_DTYPE
         return result

pandas/core/indexes/datetimelike.py

@@ -582,7 +582,8 @@ def delete(self, loc):
                 if loc.start in (0, None) or loc.stop in (len(self), None):
                     freq = self.freq
 
-        return self._shallow_copy(new_i8s, freq=freq)
+        arr = type(self._data)._simple_new(new_i8s, dtype=self.dtype, freq=freq)
+        return type(self)._simple_new(arr, name=self.name)
 
 
 class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin, Int64Index):
@@ -623,6 +624,14 @@ def _shallow_copy(self, values=None, **kwargs):
         if values is None:
             values = self._data
 
+        if isinstance(values, type(self)):
+            values = values._data
+        if isinstance(values, np.ndarray):
+            # TODO: We would rather not get here
+            if kwargs.get("freq") is not None:
+                raise ValueError(kwargs)
+            values = type(self._data)(values, dtype=self.dtype)
+
         attributes = self._get_attributes_dict()
 
         if "freq" not in kwargs and self.freq is not None:
@@ -801,7 +810,10 @@ def _union(self, other, sort):
         this, other = self._maybe_utc_convert(other)
 
         if this._can_fast_union(other):
-            return this._fast_union(other, sort=sort)
+            result = this._fast_union(other, sort=sort)
+            if result.freq is None:
+                result._set_freq("infer")
+            return result
         else:
             i8self = Int64Index._simple_new(self.asi8, name=self.name)
             i8other = Int64Index._simple_new(other.asi8, name=other.name)
@@ -934,7 +946,8 @@ def insert(self, loc, item):
             new_i8s = np.concatenate(
                 (self[:loc].asi8, [item.view(np.int64)], self[loc:].asi8)
             )
-            return self._shallow_copy(new_i8s, freq=freq)
+            arr = type(self._data)._simple_new(new_i8s, dtype=self.dtype, freq=freq)
+            return type(self)._simple_new(arr, name=self.name)
         except (AttributeError, TypeError):
 
             # fall back to object index

pandas/core/indexes/timedeltas.py

@@ -173,22 +173,15 @@ def __new__(
     def _simple_new(cls, values, name=None, freq=None, dtype=_TD_DTYPE):
         # `dtype` is passed by _shallow_copy in corner cases, should always
         #  be timedelta64[ns] if present
-
-        if not isinstance(values, TimedeltaArray):
-            values = TimedeltaArray._simple_new(values, dtype=dtype, freq=freq)
-        else:
-            if freq is None:
-                freq = values.freq
-        assert isinstance(values, TimedeltaArray), type(values)
         assert dtype == _TD_DTYPE, dtype
-        assert values.dtype == "m8[ns]", values.dtype
+        assert isinstance(values, TimedeltaArray)
+        assert freq is None or values.freq == freq
 
-        tdarr = TimedeltaArray._simple_new(values._data, freq=freq)
         result = object.__new__(cls)
-        result._data = tdarr
+        result._data = values
         result._name = name
         # For groupby perf. See note in indexes/base about _index_data
-        result._index_data = tdarr._data
+        result._index_data = values._data
 
         result._reset_identity()
         return result

pandas/core/resample.py

@@ -23,6 +23,7 @@
 from pandas.core.groupby.groupby import GroupBy, _GroupBy, _pipe_template, get_groupby
 from pandas.core.groupby.grouper import Grouper
 from pandas.core.groupby.ops import BinGrouper
+from pandas.core.indexes.api import Index
 from pandas.core.indexes.datetimes import DatetimeIndex, date_range
 from pandas.core.indexes.period import PeriodIndex, period_range
 from pandas.core.indexes.timedeltas import TimedeltaIndex, timedelta_range
@@ -424,10 +425,7 @@ def _wrap_result(self, result):
 
         if isinstance(result, ABCSeries) and result.empty:
             obj = self.obj
-            if isinstance(obj.index, PeriodIndex):
-                result.index = obj.index.asfreq(self.freq)
-            else:
-                result.index = obj.index._shallow_copy(freq=self.freq)
+            result.index = _asfreq_compat(obj.index, freq=self.freq)
             result.name = getattr(obj, "name", None)
 
         return result
@@ -1787,8 +1785,8 @@ def asfreq(obj, freq, method=None, how=None, normalize=False, fill_value=None):
 
     elif len(obj.index) == 0:
         new_obj = obj.copy()
-        new_obj.index = obj.index._shallow_copy(freq=to_offset(freq))
 
+        new_obj.index = _asfreq_compat(obj.index, freq)
     else:
         dti = date_range(obj.index[0], obj.index[-1], freq=freq)
         dti.name = obj.index.name
@@ -1797,3 +1795,28 @@ def asfreq(obj, freq, method=None, how=None, normalize=False, fill_value=None):
             new_obj.index = new_obj.index.normalize()
 
     return new_obj
+
+
+def _asfreq_compat(index, freq):
+    """
+    Helper to mimic asfreq on (empty) DatetimeIndex and TimedeltaIndex.
+
+    Parameters
+    ----------
+    index : PeriodIndex, DatetimeIndex, or TimedeltaIndex
+    freq : DateOffset
+
+    Returns
+    -------
+    same type as index
+    """
+    if len(index) != 0:
+        # This should never be reached, always checked by the caller
+        raise ValueError(
+            "Can only set arbitrary freq for empty DatetimeIndex or TimedeltaIndex"
+        )
+    if isinstance(index, PeriodIndex):
+        new_index = index.asfreq(freq=freq)
+    else:
+        new_index = Index([], dtype=index.dtype, freq=freq, name=index.name)
+    return new_index

pandas/tests/resample/test_base.py

@@ -11,6 +11,7 @@
 from pandas.core.indexes.datetimes import date_range
 from pandas.core.indexes.period import PeriodIndex, period_range
 from pandas.core.indexes.timedeltas import TimedeltaIndex, timedelta_range
+from pandas.core.resample import _asfreq_compat
 
 # a fixture value can be overridden by the test parameter value. Note that the
 # value of the fixture can be overridden this way even if the test doesn't use
@@ -103,10 +104,8 @@ def test_resample_empty_series(freq, empty_series, resample_method):
     result = getattr(s.resample(freq), resample_method)()
 
     expected = s.copy()
-    if isinstance(s.index, PeriodIndex):
-        expected.index = s.index.asfreq(freq=freq)
-    else:
-        expected.index = s.index._shallow_copy(freq=freq)
+    expected.index = _asfreq_compat(s.index, freq)
+
     tm.assert_index_equal(result.index, expected.index)
     assert result.index.freq == expected.index.freq
     tm.assert_series_equal(result, expected, check_dtype=False)
@@ -119,10 +118,8 @@ def test_resample_count_empty_series(freq, empty_series, resample_method):
     # GH28427
     result = getattr(empty_series.resample(freq), resample_method)()
 
-    if isinstance(empty_series.index, PeriodIndex):
-        index = empty_series.index.asfreq(freq=freq)
-    else:
-        index = empty_series.index._shallow_copy(freq=freq)
+    index = _asfreq_compat(empty_series.index, freq)
+
     expected = pd.Series([], dtype="int64", index=index, name=empty_series.name)
 
     tm.assert_series_equal(result, expected)
@@ -141,10 +138,8 @@ def test_resample_empty_dataframe(empty_frame, freq, resample_method):
         # GH14962
         expected = Series([], dtype=object)
 
-    if isinstance(df.index, PeriodIndex):
-        expected.index = df.index.asfreq(freq=freq)
-    else:
-        expected.index = df.index._shallow_copy(freq=freq)
+    expected.index = _asfreq_compat(df.index, freq)
+
     tm.assert_index_equal(result.index, expected.index)
     assert result.index.freq == expected.index.freq
     tm.assert_almost_equal(result, expected, check_dtype=False)
@@ -162,10 +157,8 @@ def test_resample_count_empty_dataframe(freq, empty_frame):
 
     result = empty_frame.resample(freq).count()
 
-    if isinstance(empty_frame.index, PeriodIndex):
-        index = empty_frame.index.asfreq(freq=freq)
-    else:
-        index = empty_frame.index._shallow_copy(freq=freq)
+    index = _asfreq_compat(empty_frame.index, freq)
+
     expected = pd.DataFrame({"a": []}, dtype="int64", index=index)
 
     tm.assert_frame_equal(result, expected)
@@ -181,10 +174,8 @@ def test_resample_size_empty_dataframe(freq, empty_frame):
 
     result = empty_frame.resample(freq).size()
 
-    if isinstance(empty_frame.index, PeriodIndex):
-        index = empty_frame.index.asfreq(freq=freq)
-    else:
-        index = empty_frame.index._shallow_copy(freq=freq)
+    index = _asfreq_compat(empty_frame.index, freq)
+
     expected = pd.Series([], dtype="int64", index=index)
 
     tm.assert_series_equal(result, expected)