Skip to content

Commit c51867a

Browse files
authored
REF: share ExtensionIndex.insert-> Index.insert (#44170)
1 parent aced6ee commit c51867a

File tree

4 files changed

+24
-29
lines changed

4 files changed

+24
-29
lines changed

pandas/core/dtypes/missing.py

+4
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
needs_i8_conversion,
3838
)
3939
from pandas.core.dtypes.dtypes import (
40+
CategoricalDtype,
4041
ExtensionDtype,
4142
IntervalDtype,
4243
PeriodDtype,
@@ -641,5 +642,8 @@ def is_valid_na_for_dtype(obj, dtype: DtypeObj) -> bool:
641642
elif isinstance(dtype, IntervalDtype):
642643
return lib.is_float(obj) or obj is None or obj is libmissing.NA
643644

645+
elif isinstance(dtype, CategoricalDtype):
646+
return is_valid_na_for_dtype(obj, dtype.categories.dtype)
647+
644648
# fallback, default to allowing NaN, None, NA, NaT
645649
return not isinstance(obj, (np.datetime64, np.timedelta64, Decimal))

pandas/core/indexes/base.py

+11-4
Original file line numberDiff line numberDiff line change
@@ -6432,14 +6432,21 @@ def insert(self, loc: int, item) -> Index:
64326432
if is_valid_na_for_dtype(item, self.dtype) and self.dtype != object:
64336433
item = self._na_value
64346434

6435+
arr = self._values
6436+
64356437
try:
6436-
item = self._validate_fill_value(item)
6437-
except TypeError:
6438+
if isinstance(arr, ExtensionArray):
6439+
res_values = arr.insert(loc, item)
6440+
return type(self)._simple_new(res_values, name=self.name)
6441+
else:
6442+
item = self._validate_fill_value(item)
6443+
except (TypeError, ValueError):
6444+
# e.g. trying to insert an integer into a DatetimeIndex
6445+
# We cannot keep the same dtype, so cast to the (often object)
6446+
# minimal shared dtype before doing the insert.
64386447
dtype = self._find_common_type_compat(item)
64396448
return self.astype(dtype).insert(loc, item)
64406449

6441-
arr = self._values
6442-
64436450
if arr.dtype != object or not isinstance(
64446451
item, (tuple, np.datetime64, np.timedelta64)
64456452
):

pandas/core/indexes/extension.py

-25
Original file line numberDiff line numberDiff line change
@@ -134,31 +134,6 @@ class ExtensionIndex(Index):
134134

135135
# ---------------------------------------------------------------------
136136

137-
def insert(self, loc: int, item) -> Index:
138-
"""
139-
Make new Index inserting new item at location. Follows
140-
Python list.append semantics for negative values.
141-
142-
Parameters
143-
----------
144-
loc : int
145-
item : object
146-
147-
Returns
148-
-------
149-
new_index : Index
150-
"""
151-
try:
152-
result = self._data.insert(loc, item)
153-
except (ValueError, TypeError):
154-
# e.g. trying to insert an integer into a DatetimeIndex
155-
# We cannot keep the same dtype, so cast to the (often object)
156-
# minimal shared dtype before doing the insert.
157-
dtype = self._find_common_type_compat(item)
158-
return self.astype(dtype).insert(loc, item)
159-
else:
160-
return type(self)._simple_new(result, name=self.name)
161-
162137
def _validate_fill_value(self, value):
163138
"""
164139
Convert value to be insertable to underlying array.

pandas/tests/dtypes/test_missing.py

+9
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
is_scalar,
1919
)
2020
from pandas.core.dtypes.dtypes import (
21+
CategoricalDtype,
2122
DatetimeTZDtype,
2223
IntervalDtype,
2324
PeriodDtype,
@@ -739,3 +740,11 @@ def test_is_valid_na_for_dtype_interval(self):
739740

740741
dtype = IntervalDtype("datetime64[ns]", "both")
741742
assert not is_valid_na_for_dtype(NaT, dtype)
743+
744+
def test_is_valid_na_for_dtype_categorical(self):
745+
dtype = CategoricalDtype(categories=[0, 1, 2])
746+
assert is_valid_na_for_dtype(np.nan, dtype)
747+
748+
assert not is_valid_na_for_dtype(NaT, dtype)
749+
assert not is_valid_na_for_dtype(np.datetime64("NaT", "ns"), dtype)
750+
assert not is_valid_na_for_dtype(np.timedelta64("NaT", "ns"), dtype)

0 commit comments

Comments
 (0)