Shift to adding frequency to shallow copy of DatetimeIndex instead of modifying take function

Author: annika-rudolph

pandas/core/arrays/datetimelike.py

@@ -66,7 +66,6 @@
     ScalarIndexer,
     Self,
     SequenceIndexer,
-    TakeIndexer,
     TimeAmbiguous,
     TimeNonexistent,
     npt,
@@ -2360,27 +2359,6 @@ def interpolate(
             return self
         return type(self)._simple_new(out_data, dtype=self.dtype)
 
-    def take(
-        self,
-        indices: TakeIndexer,
-        *,
-        allow_fill: bool = False,
-        fill_value: Any = None,
-        axis: AxisInt = 0,
-    ) -> Self:
-        result = super().take(
-            indices=indices, allow_fill=allow_fill, fill_value=fill_value, axis=axis
-        )
-
-        indices = np.asarray(indices, dtype=np.intp)
-        maybe_slice = lib.maybe_indices_to_slice(indices, len(self))
-
-        if isinstance(maybe_slice, slice):
-            freq = self._get_getitem_freq(maybe_slice)
-            result.freq = freq
-
-        return result
-
     # --------------------------------------------------------------
     # Unsorted
 

pandas/core/indexes/base.py

@@ -748,7 +748,7 @@ def _format_duplicate_message(self) -> DataFrame:
     # --------------------------------------------------------------------
     # Index Internals Methods
 
-    def _shallow_copy(self, values, name: Hashable = no_default) -> Self:
+    def _shallow_copy(self, values, name: Hashable = no_default, **kwargs) -> Self:
         """
         Create a new Index with the same class as the caller, don't copy the
         data, use the same object attributes with passed in attributes taking

pandas/core/indexes/datetimelike.py

@@ -22,6 +22,7 @@
     Timedelta,
     lib,
 )
+from pandas._libs.lib import no_default
 from pandas._libs.tslibs import (
     BaseOffset,
     Resolution,
@@ -68,7 +69,10 @@
 from pandas.core.tools.timedeltas import to_timedelta
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import (
+        Hashable,
+        Sequence,
+    )
     from datetime import datetime
 
     from pandas._typing import (
@@ -838,3 +842,20 @@ def take(
             freq = self._data._get_getitem_freq(maybe_slice)
             result._data._freq = freq
         return result
+
+    @doc(Index._shallow_copy)
+    def _shallow_copy( # type: ignore[override]
+        self,
+        values,
+        name: Hashable = no_default,
+        level_codes=no_default,
+    ) -> Self:
+        name = self._name if name is no_default else name
+        result = self._simple_new(values, name=name, refs=self._references)
+        if level_codes is not no_default:
+            indices = np.asarray(level_codes, dtype=np.intp)
+            maybe_slice = lib.maybe_indices_to_slice(indices, len(values))
+            if isinstance(maybe_slice, slice):
+                freq = self._data._get_getitem_freq(maybe_slice)
+                result._data._freq = freq
+        return result

pandas/core/indexes/multi.py

@@ -1244,7 +1244,9 @@ def _constructor(self) -> Callable[..., MultiIndex]:  # type: ignore[override]
         return type(self).from_tuples
 
     @doc(Index._shallow_copy)
-    def _shallow_copy(self, values: np.ndarray, name=lib.no_default) -> MultiIndex:
+    def _shallow_copy(
+        self, values: np.ndarray, name=lib.no_default, **kwargs
+    ) -> MultiIndex:
         names = name if name is not lib.no_default else self.names
 
         return type(self).from_tuples(values, sortorder=None, names=names)
@@ -1714,7 +1716,7 @@ def _get_level_values(self, level: int, unique: bool = False) -> Index:
         if unique:
             level_codes = algos.unique(level_codes)
         filled = algos.take_nd(lev._values, level_codes, fill_value=lev._na_value)
-        return lev._shallow_copy(filled, name=name)
+        return lev._shallow_copy(filled, name=name, level_codes=level_codes)
 
     # error: Signature of "get_level_values" incompatible with supertype "Index"
     def get_level_values(self, level) -> Index:  # type: ignore[override]

pandas/core/indexes/range.py

@@ -472,7 +472,7 @@ def __iter__(self) -> Iterator[int]:
         yield from self._range
 
     @doc(Index._shallow_copy)
-    def _shallow_copy(self, values, name: Hashable = no_default):
+    def _shallow_copy(self, values, name: Hashable = no_default, **kwargs):
         name = self._name if name is no_default else name
 
         if values.dtype.kind == "f":

pandas/tests/indexes/multi/test_get_level_values.py

@@ -126,9 +126,7 @@ def test_values_loses_freq_of_underlying_index():
 
 def test_get_level_values_gets_frequency_correctly():
     # GH#57949 GH#58327
-    datetime_index = pd.date_range(
-        start=pd.to_datetime("1/1/2018"), periods=4, freq="YS"
-    )
+    datetime_index = date_range(start=pd.to_datetime("1/1/2018"), periods=4, freq="YS")
     other_index = ["A"]
     multi_index = MultiIndex.from_product([datetime_index, other_index])