fixup insert

Author: MarcoGorelli

pandas-stubs/core/arrays/categorical.pyi

@@ -110,6 +110,7 @@ class Categorical(ExtensionArray):
     def take(
         self, indexer: TakeIndexer, *, allow_fill: bool = ..., fill_value=...
     ) -> Categorical: ...
+    def __contains__(self, key) -> bool: ...
     @overload
     def __getitem__(self, key: ScalarIndexer) -> Any: ...
     @overload
@@ -120,6 +121,7 @@ class Categorical(ExtensionArray):
     def __setitem__(self, key, value) -> None: ...
     def min(self, *, skipna: bool = ...): ...
     def max(self, *, skipna: bool = ...): ...
+    def unique(self): ...
     def equals(self, other): ...
     def describe(self): ...
     def repeat(self, repeats, axis=...): ...

pandas-stubs/core/indexes/base.pyi

@@ -524,6 +524,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
         self, start: SliceType = None, end: SliceType = None, step: int | None = None
     ): ...
     def delete(self, loc: int | AnyArrayLikeInt | Sequence[int]) -> Self: ...
+    @overload
+    def insert(self, loc: int, item: S1) -> Self: ...
+    @overload
     def insert(self, loc: int, item: object) -> Index: ...
     def drop(self, labels, errors: IgnoreRaise = "raise") -> Self: ...
     @property

pandas-stubs/core/indexes/category.pyi

@@ -37,3 +37,6 @@ class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate):
     def is_monotonic_increasing(self) -> bool: ...
     @property
     def is_monotonic_decreasing(self) -> bool: ...
+    # `item` might be `S1` but not one of the categories, thus changing
+    # the return type from `CategoricalIndex` to `Index`.
+    def insert(self, loc: int, item: object) -> Index: ...  # type: ignore[override]

tests/indexes/test_categoricalindex.py

@@ -22,9 +22,10 @@ def test_categoricalindex_unique() -> None:
 
 def test_categoricalindex_reindex() -> None:
     ci = pd.CategoricalIndex(["a", "b"])
-    reindexed = ci.reindex(["b", "c"])
-    check(assert_type(reindexed[0], pd.Index), pd.Index)
-    check(assert_type(reindexed[1], np_1darray[np.intp] | None), np_1darray)
+    check(
+        assert_type(ci.reindex([0, 1]), tuple[pd.Index, np_1darray[np.intp] | None]),
+        tuple,
+    )
 
 
 def test_categoricalindex_delete() -> None: