Skip to content

Commit

Permalink
Shift to adding frequency to shallow copy of DatetimeIndex instead of…
Browse files Browse the repository at this point in the history
… modifying take function
  • Loading branch information
annika-rudolph committed Apr 25, 2024
1 parent 9e2c11d commit 5440815
Show file tree
Hide file tree
Showing 6 changed files with 29 additions and 30 deletions.
22 changes: 0 additions & 22 deletions pandas/core/arrays/datetimelike.py
Expand Up @@ -66,7 +66,6 @@
ScalarIndexer,
Self,
SequenceIndexer,
TakeIndexer,
TimeAmbiguous,
TimeNonexistent,
npt,
Expand Down Expand Up @@ -2355,27 +2354,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

Expand Down
2 changes: 1 addition & 1 deletion pandas/core/indexes/base.py
Expand Up @@ -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
Expand Down
23 changes: 22 additions & 1 deletion pandas/core/indexes/datetimelike.py
Expand Up @@ -22,6 +22,7 @@
Timedelta,
lib,
)
from pandas._libs.lib import no_default
from pandas._libs.tslibs import (
BaseOffset,
Resolution,
Expand Down Expand Up @@ -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 (
Expand Down Expand Up @@ -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
6 changes: 4 additions & 2 deletions pandas/core/indexes/multi.py
Expand Up @@ -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)
Expand Down Expand Up @@ -1713,7 +1715,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]
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/indexes/range.py
Expand Up @@ -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":
Expand Down
4 changes: 1 addition & 3 deletions pandas/tests/indexes/multi/test_get_level_values.py
Expand Up @@ -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])

Expand Down

0 comments on commit 5440815

Please sign in to comment.