Skip to content

Commit f8e905e

Browse files
jbrockmendeljreback
authored andcommitted
REF: share searchsorted between DTI/TDI/PI, insert between DTI/TDI (#31143)
1 parent f4de727 commit f8e905e

File tree

4 files changed

+73
-150
lines changed

4 files changed

+73
-150
lines changed

pandas/core/indexes/datetimelike.py

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,13 @@
2727
)
2828
from pandas.core.dtypes.concat import concat_compat
2929
from pandas.core.dtypes.generic import ABCIndex, ABCIndexClass, ABCSeries
30-
from pandas.core.dtypes.missing import isna
30+
from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna
3131

3232
from pandas.core import algorithms
3333
from pandas.core.accessor import PandasDelegate
3434
from pandas.core.arrays import DatetimeArray, ExtensionArray, TimedeltaArray
3535
from pandas.core.arrays.datetimelike import DatetimeLikeArrayMixin
36+
from pandas.core.base import _shared_docs
3637
import pandas.core.indexes.base as ibase
3738
from pandas.core.indexes.base import Index, _index_shared_docs
3839
from pandas.core.indexes.extension import (
@@ -222,6 +223,18 @@ def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
222223
self, indices, axis, allow_fill, fill_value, **kwargs
223224
)
224225

226+
@Appender(_shared_docs["searchsorted"])
227+
def searchsorted(self, value, side="left", sorter=None):
228+
if isinstance(value, str):
229+
raise TypeError(
230+
"searchsorted requires compatible dtype or scalar, "
231+
f"not {type(value).__name__}"
232+
)
233+
if isinstance(value, Index):
234+
value = value._data
235+
236+
return self._data.searchsorted(value, side=side, sorter=sorter)
237+
225238
_can_hold_na = True
226239

227240
_na_value = NaT
@@ -883,6 +896,60 @@ def _wrap_joined_index(self, joined, other):
883896
kwargs["tz"] = getattr(other, "tz", None)
884897
return self._simple_new(joined, name, **kwargs)
885898

899+
# --------------------------------------------------------------------
900+
# List-Like Methods
901+
902+
def insert(self, loc, item):
903+
"""
904+
Make new Index inserting new item at location
905+
Parameters
906+
----------
907+
loc : int
908+
item : object
909+
if not either a Python datetime or a numpy integer-like, returned
910+
Index dtype will be object rather than datetime.
911+
Returns
912+
-------
913+
new_index : Index
914+
"""
915+
if isinstance(item, self._data._recognized_scalars):
916+
item = self._data._scalar_type(item)
917+
elif is_valid_nat_for_dtype(item, self.dtype):
918+
# GH 18295
919+
item = self._na_value
920+
elif is_scalar(item) and isna(item):
921+
raise TypeError(
922+
f"cannot insert {type(self).__name__} with incompatible label"
923+
)
924+
925+
freq = None
926+
if isinstance(item, self._data._scalar_type) or item is NaT:
927+
self._data._check_compatible_with(item, setitem=True)
928+
929+
# check freq can be preserved on edge cases
930+
if self.size and self.freq is not None:
931+
if item is NaT:
932+
pass
933+
elif (loc == 0 or loc == -len(self)) and item + self.freq == self[0]:
934+
freq = self.freq
935+
elif (loc == len(self)) and item - self.freq == self[-1]:
936+
freq = self.freq
937+
item = item.asm8
938+
939+
try:
940+
new_i8s = np.concatenate(
941+
(self[:loc].asi8, [item.view(np.int64)], self[loc:].asi8)
942+
)
943+
return self._shallow_copy(new_i8s, freq=freq)
944+
except (AttributeError, TypeError):
945+
946+
# fall back to object index
947+
if isinstance(item, str):
948+
return self.astype(object).insert(loc, item)
949+
raise TypeError(
950+
f"cannot insert {type(self).__name__} with incompatible label"
951+
)
952+
886953

887954
class DatetimelikeDelegateMixin(PandasDelegate):
888955
"""

pandas/core/indexes/datetimes.py

Lines changed: 2 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,18 @@
1414
tslib as libts,
1515
)
1616
from pandas._libs.tslibs import ccalendar, fields, parsing, timezones
17-
from pandas.util._decorators import Appender, Substitution, cache_readonly
17+
from pandas.util._decorators import cache_readonly
1818

1919
from pandas.core.dtypes.common import _NS_DTYPE, is_float, is_integer, is_scalar
2020
from pandas.core.dtypes.dtypes import DatetimeTZDtype
21-
from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna
21+
from pandas.core.dtypes.missing import isna
2222

2323
from pandas.core.accessor import delegate_names
2424
from pandas.core.arrays.datetimes import (
2525
DatetimeArray,
2626
tz_to_dtype,
2727
validate_tz_from_dtype,
2828
)
29-
from pandas.core.base import _shared_docs
3029
import pandas.core.common as com
3130
from pandas.core.indexes.base import Index, maybe_extract_name
3231
from pandas.core.indexes.datetimelike import (
@@ -825,19 +824,6 @@ def slice_indexer(self, start=None, end=None, step=None, kind=None):
825824

826825
# --------------------------------------------------------------------
827826

828-
@Substitution(klass="DatetimeIndex")
829-
@Appender(_shared_docs["searchsorted"])
830-
def searchsorted(self, value, side="left", sorter=None):
831-
if isinstance(value, str):
832-
raise TypeError(
833-
"searchsorted requires compatible dtype or scalar, "
834-
f"not {type(value).__name__}"
835-
)
836-
if isinstance(value, Index):
837-
value = value._data
838-
839-
return self._data.searchsorted(value, side=side)
840-
841827
def is_type_compatible(self, typ) -> bool:
842828
return typ == self.inferred_type or typ == "datetime"
843829

@@ -847,60 +833,6 @@ def inferred_type(self) -> str:
847833
# sure we can't have ambiguous indexing
848834
return "datetime64"
849835

850-
def insert(self, loc, item):
851-
"""
852-
Make new Index inserting new item at location
853-
854-
Parameters
855-
----------
856-
loc : int
857-
item : object
858-
if not either a Python datetime or a numpy integer-like, returned
859-
Index dtype will be object rather than datetime.
860-
861-
Returns
862-
-------
863-
new_index : Index
864-
"""
865-
if isinstance(item, self._data._recognized_scalars):
866-
item = self._data._scalar_type(item)
867-
elif is_valid_nat_for_dtype(item, self.dtype):
868-
# GH 18295
869-
item = self._na_value
870-
elif is_scalar(item) and isna(item):
871-
# i.e. timedeltat64("NaT")
872-
raise TypeError(
873-
f"cannot insert {type(self).__name__} with incompatible label"
874-
)
875-
876-
freq = None
877-
if isinstance(item, self._data._scalar_type) or item is NaT:
878-
self._data._check_compatible_with(item, setitem=True)
879-
880-
# check freq can be preserved on edge cases
881-
if self.size and self.freq is not None:
882-
if item is NaT:
883-
pass
884-
elif (loc == 0 or loc == -len(self)) and item + self.freq == self[0]:
885-
freq = self.freq
886-
elif (loc == len(self)) and item - self.freq == self[-1]:
887-
freq = self.freq
888-
item = item.asm8
889-
890-
try:
891-
new_i8s = np.concatenate(
892-
(self[:loc].asi8, [item.view(np.int64)], self[loc:].asi8)
893-
)
894-
return self._shallow_copy(new_i8s, freq=freq)
895-
except (AttributeError, TypeError):
896-
897-
# fall back to object index
898-
if isinstance(item, str):
899-
return self.astype(object).insert(loc, item)
900-
raise TypeError(
901-
f"cannot insert {type(self).__name__} with incompatible label"
902-
)
903-
904836
def indexer_at_time(self, time, asof=False):
905837
"""
906838
Return index locations of index values at particular time of day

pandas/core/indexes/period.py

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from pandas._libs.tslibs import NaT, frequencies as libfrequencies, resolution
88
from pandas._libs.tslibs.parsing import parse_time_string
99
from pandas._libs.tslibs.period import Period
10-
from pandas.util._decorators import Appender, Substitution, cache_readonly
10+
from pandas.util._decorators import Appender, cache_readonly
1111

1212
from pandas.core.dtypes.common import (
1313
ensure_platform_int,
@@ -29,7 +29,6 @@
2929
raise_on_incompatible,
3030
validate_dtype_freq,
3131
)
32-
from pandas.core.base import _shared_docs
3332
import pandas.core.common as com
3433
import pandas.core.indexes.base as ibase
3534
from pandas.core.indexes.base import (
@@ -455,11 +454,6 @@ def astype(self, dtype, copy=True, how="start"):
455454
# TODO: should probably raise on `how` here, so we don't ignore it.
456455
return super().astype(dtype, copy=copy)
457456

458-
@Substitution(klass="PeriodIndex")
459-
@Appender(_shared_docs["searchsorted"])
460-
def searchsorted(self, value, side="left", sorter=None):
461-
return self._data.searchsorted(value, side=side, sorter=sorter)
462-
463457
@property
464458
def is_full(self) -> bool:
465459
"""

pandas/core/indexes/timedeltas.py

Lines changed: 2 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import numpy as np
55

66
from pandas._libs import NaT, Timedelta, index as libindex
7-
from pandas.util._decorators import Appender, Substitution
7+
from pandas.util._decorators import Appender
88

99
from pandas.core.dtypes.common import (
1010
_TD_DTYPE,
@@ -16,12 +16,11 @@
1616
is_timedelta64_ns_dtype,
1717
pandas_dtype,
1818
)
19-
from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna
19+
from pandas.core.dtypes.missing import isna
2020

2121
from pandas.core.accessor import delegate_names
2222
from pandas.core.arrays import datetimelike as dtl
2323
from pandas.core.arrays.timedeltas import TimedeltaArray, _is_convertible_to_td
24-
from pandas.core.base import _shared_docs
2524
import pandas.core.common as com
2625
from pandas.core.indexes.base import Index, _index_shared_docs, maybe_extract_name
2726
from pandas.core.indexes.datetimelike import (
@@ -345,82 +344,13 @@ def _partial_td_slice(self, key):
345344

346345
raise NotImplementedError
347346

348-
@Substitution(klass="TimedeltaIndex")
349-
@Appender(_shared_docs["searchsorted"])
350-
def searchsorted(self, value, side="left", sorter=None):
351-
if isinstance(value, str):
352-
raise TypeError(
353-
"searchsorted requires compatible dtype or scalar, "
354-
f"not {type(value).__name__}"
355-
)
356-
if isinstance(value, Index):
357-
value = value._data
358-
359-
return self._data.searchsorted(value, side=side, sorter=sorter)
360-
361347
def is_type_compatible(self, typ) -> bool:
362348
return typ == self.inferred_type or typ == "timedelta"
363349

364350
@property
365351
def inferred_type(self) -> str:
366352
return "timedelta64"
367353

368-
def insert(self, loc, item):
369-
"""
370-
Make new Index inserting new item at location
371-
372-
Parameters
373-
----------
374-
loc : int
375-
item : object
376-
If not either a Python datetime or a numpy integer-like, returned
377-
Index dtype will be object rather than datetime.
378-
379-
Returns
380-
-------
381-
new_index : Index
382-
"""
383-
# try to convert if possible
384-
if isinstance(item, self._data._recognized_scalars):
385-
item = self._data._scalar_type(item)
386-
elif is_valid_nat_for_dtype(item, self.dtype):
387-
# GH 18295
388-
item = self._na_value
389-
elif is_scalar(item) and isna(item):
390-
# i.e. datetime64("NaT")
391-
raise TypeError(
392-
f"cannot insert {type(self).__name__} with incompatible label"
393-
)
394-
395-
freq = None
396-
if isinstance(item, self._data._scalar_type) or item is NaT:
397-
self._data._check_compatible_with(item, setitem=True)
398-
399-
# check freq can be preserved on edge cases
400-
if self.size and self.freq is not None:
401-
if item is NaT:
402-
pass
403-
elif (loc == 0 or loc == -len(self)) and item + self.freq == self[0]:
404-
freq = self.freq
405-
elif (loc == len(self)) and item - self.freq == self[-1]:
406-
freq = self.freq
407-
item = item.asm8
408-
409-
try:
410-
new_i8s = np.concatenate(
411-
(self[:loc].asi8, [item.view(np.int64)], self[loc:].asi8)
412-
)
413-
tda = type(self._data)._simple_new(new_i8s, freq=freq)
414-
return self._shallow_copy(tda)
415-
except (AttributeError, TypeError):
416-
417-
# fall back to object index
418-
if isinstance(item, str):
419-
return self.astype(object).insert(loc, item)
420-
raise TypeError(
421-
f"cannot insert {type(self).__name__} with incompatible label"
422-
)
423-
424354

425355
TimedeltaIndex._add_logical_methods_disabled()
426356

0 commit comments

Comments
 (0)