BUG: TDI/DTI _shallow_copy creating invalid arrays (pandas-dev#30764)

Author: jbrockmendel

pandas/core/indexes/datetimelike.py

@@ -28,7 +28,12 @@
 
 from pandas.core import algorithms
 from pandas.core.accessor import PandasDelegate
-from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
+from pandas.core.arrays import (
+    DatetimeArray,
+    ExtensionArray,
+    ExtensionOpsMixin,
+    TimedeltaArray,
+)
 from pandas.core.arrays.datetimelike import (
     DatetimeLikeArrayMixin,
     _ensure_datetimelike_to_i8,
@@ -251,15 +256,10 @@ def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
         if isinstance(maybe_slice, slice):
             return self[maybe_slice]
 
-        taken = ExtensionIndex.take(
+        return ExtensionIndex.take(
             self, indices, axis, allow_fill, fill_value, **kwargs
         )
 
-        # keep freq in PeriodArray/Index, reset otherwise
-        freq = self.freq if is_period_dtype(self) else None
-        assert taken.freq == freq, (taken.freq, freq, taken)
-        return self._shallow_copy(taken, freq=freq)
-
     _can_hold_na = True
 
     _na_value = NaT
@@ -486,8 +486,8 @@ def isin(self, values, level=None):
     @Appender(_index_shared_docs["repeat"] % _index_doc_kwargs)
     def repeat(self, repeats, axis=None):
         nv.validate_repeat(tuple(), dict(axis=axis))
-        freq = self.freq if is_period_dtype(self) else None
-        return self._shallow_copy(self.asi8.repeat(repeats), freq=freq)
+        result = type(self._data)(self.asi8.repeat(repeats), dtype=self.dtype)
+        return self._shallow_copy(result)
 
     @Appender(_index_shared_docs["where"] % _index_doc_kwargs)
     def where(self, cond, other=None):
@@ -650,6 +650,22 @@ def _set_freq(self, freq):
 
         self._data._freq = freq
 
+    def _shallow_copy(self, values=None, **kwargs):
+        if values is None:
+            values = self._data
+        if isinstance(values, type(self)):
+            values = values._data
+
+        attributes = self._get_attributes_dict()
+
+        if "freq" not in kwargs and self.freq is not None:
+            if isinstance(values, (DatetimeArray, TimedeltaArray)):
+                if values.freq is None:
+                    del attributes["freq"]
+
+        attributes.update(kwargs)
+        return self._simple_new(values, **attributes)
+
     # --------------------------------------------------------------------
     # Set Operation Methods
 

pandas/tests/indexes/datetimes/test_constructors.py

@@ -29,6 +29,25 @@ def test_freq_validation_with_nat(self, dt_cls):
         with pytest.raises(ValueError, match=msg):
             dt_cls([pd.NaT, pd.Timestamp("2011-01-01").value], freq="D")
 
+    # TODO: better place for tests shared by DTI/TDI?
+    @pytest.mark.parametrize(
+        "index",
+        [
+            pd.date_range("2016-01-01", periods=5, tz="US/Pacific"),
+            pd.timedelta_range("1 Day", periods=5),
+        ],
+    )
+    def test_shallow_copy_inherits_array_freq(self, index):
+        # If we pass a DTA/TDA to shallow_copy and dont specify a freq,
+        #  we should inherit the array's freq, not our own.
+        array = index._data
+
+        arr = array[[0, 3, 2, 4, 1]]
+        assert arr.freq is None
+
+        result = index._shallow_copy(arr)
+        assert result.freq is None
+
     def test_categorical_preserves_tz(self):
         # GH#18664 retain tz when going DTI-->Categorical-->DTI
         # TODO: parametrize over DatetimeIndex/DatetimeArray