Type CategoricalIndex (#1459)

* wip

* bit more

* testing

* more tests

* fixup `insert`

* include `np.integer`, as per upstream

* type `categories` too

* update pyproject.toml

Author: MarcoGorelli

pandas-stubs/core/indexes/base.pyi

@@ -69,6 +69,7 @@ from pandas._typing import (
     T_COMPLEX,
     AnyAll,
     AnyArrayLike,
+    AnyArrayLikeInt,
     ArrayLike,
     AxesData,
     CategoryDtypeArg,
@@ -406,7 +407,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     notnull = ...
     def fillna(self, value=...): ...
     def dropna(self, how: AnyAll = "any") -> Self: ...
-    def unique(self, level=...) -> Self: ...
+    def unique(self, level: Hashable | None = None) -> Self: ...
     def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ...
     def duplicated(self, keep: DropKeep = "first") -> np_1darray[np.bool]: ...
     def __and__(self, other: Never) -> Never: ...
@@ -442,12 +443,12 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     ) -> np_1darray[np.intp]: ...
     def reindex(
         self,
-        target,
-        method: ReindexMethod | None = ...,
-        level=...,
-        limit=...,
-        tolerance=...,
-    ): ...
+        target: Iterable[Any],
+        method: ReindexMethod | None = None,
+        level: int | None = None,
+        limit: int | None = None,
+        tolerance: Scalar | AnyArrayLike | Sequence[Scalar] | None = None,
+    ) -> tuple[Index, np_1darray[np.intp] | None]: ...
     @overload
     def join(
         self,
@@ -483,7 +484,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
         cond: Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool],
         other: Scalar | AnyArrayLike | None = None,
     ) -> Index: ...
-    def __contains__(self, key) -> bool: ...
+    def __contains__(self, key: Hashable) -> bool: ...
     @final
     def __setitem__(self, key, value) -> None: ...
     @overload
@@ -500,7 +501,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     @overload
     def append(self, other: Index | Sequence[Index]) -> Index: ...
     def putmask(self, mask, value): ...
-    def equals(self, other) -> bool: ...
+    def equals(self, other: object) -> bool: ...
     @final
     def identical(self, other) -> bool: ...
     @final
@@ -534,8 +535,13 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     def slice_locs(
         self, start: SliceType = None, end: SliceType = None, step: int | None = None
     ): ...
-    def delete(self, loc) -> Self: ...
-    def insert(self, loc, item) -> Self: ...
+    def delete(
+        self, loc: np.integer | 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
     def shape(self) -> tuple[int, ...]: ...

pandas-stubs/core/indexes/category.pyi

@@ -2,7 +2,6 @@ from collections.abc import (
     Hashable,
     Iterable,
 )
-from typing import final
 
 import numpy as np
 from pandas.core import accessor
@@ -11,7 +10,11 @@ from pandas.core.indexes.base import Index
 from pandas.core.indexes.extension import ExtensionIndex
 from typing_extensions import Self
 
-from pandas._typing import S1
+from pandas._typing import (
+    S1,
+    Dtype,
+    ListLike,
+)
 
 class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate):
     codes: np.ndarray = ...
@@ -21,28 +24,20 @@ class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate):
     def __new__(
         cls,
         data: Iterable[S1] = ...,
-        categories=...,
-        ordered=...,
-        dtype=...,
-        copy: bool = ...,
-        name: Hashable = ...,
+        categories: ListLike | None = None,
+        ordered: bool | None = None,
+        dtype: Dtype | None = None,
+        copy: bool = False,
+        name: Hashable | None = None,
     ) -> Self: ...
-    def equals(self, other): ...
     @property
     def inferred_type(self) -> str: ...
     @property
-    def values(self): ...
-    def __contains__(self, key) -> bool: ...
-    @property
     def is_unique(self) -> bool: ...
     @property
     def is_monotonic_increasing(self) -> bool: ...
     @property
     def is_monotonic_decreasing(self) -> bool: ...
-    def unique(self, level=...): ...
-    def reindex(self, target, method=..., level=..., limit=..., tolerance=...): ...
-    @final
-    def get_indexer(self, target, method=..., limit=..., tolerance=...): ...
-    def get_indexer_non_unique(self, target): ...
-    def delete(self, loc): ...
-    def insert(self, loc, item): ...
+    # `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]

pandas-stubs/core/indexes/range.pyi

@@ -61,7 +61,6 @@ class RangeIndex(_IndexSubclassBase[int, np.int64]):
     def is_monotonic_decreasing(self) -> bool: ...
     @property
     def has_duplicates(self) -> bool: ...
-    def __contains__(self, key: int | np.integer) -> bool: ...
     def factorize(
         self, sort: bool = False, use_na_sentinel: bool = True
     ) -> tuple[np_1darray[np.intp], RangeIndex]: ...

pyproject.toml

@@ -204,10 +204,6 @@ ignore = [
   "PYI042",  # https://docs.astral.sh/ruff/rules/snake-case-type-alias/
   "ERA001", "PLR0402", "PLC0105"
 ]
-"*category.pyi" = [
-  # TODO: remove when pandas-dev/pandas-stubs#1443 is resolved
-  "ANN001", "ANN201", "ANN204", "ANN206",
-]
 "*series.pyi" = [
   # TODO: remove when pandas-dev/pandas-stubs#1444 is resolved
   "ANN001", "ANN201", "ANN204", "ANN206",

tests/indexes/test_categoricalindex.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from typing_extensions import (
+    assert_type,
+)
+
+from tests import (
+    check,
+    np_1darray,
+)
+
+
+def test_categoricalindex_unique() -> None:
+    ci = pd.CategoricalIndex(["a", "b"])
+    check(
+        assert_type(ci.unique(), "pd.CategoricalIndex[str]"),
+        pd.CategoricalIndex,
+    )
+
+
+def test_categoricalindex_reindex() -> None:
+    ci = pd.CategoricalIndex(["a", "b"])
+    check(
+        assert_type(ci.reindex([0, 1]), tuple[pd.Index, np_1darray[np.intp] | None]),
+        tuple,
+    )
+
+
+def test_categoricalindex_delete() -> None:
+    ci = pd.CategoricalIndex(["a", "b"])
+    check(assert_type(ci.delete(0), "pd.CategoricalIndex[str]"), pd.CategoricalIndex)
+    check(
+        assert_type(ci.delete([0, 1]), "pd.CategoricalIndex[str]"), pd.CategoricalIndex
+    )
+
+
+def test_categoricalindex_insert() -> None:
+    ci = pd.CategoricalIndex(["a", "b"])
+    check(assert_type(ci.insert(0, "c"), pd.Index), pd.Index)