diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index c4bf83138..7f9f80dc4 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -69,6 +69,7 @@ from pandas._typing import ( T_COMPLEX, AnyAll, AnyArrayLike, + AnyArrayLikeInt, ArrayLike, AxesData, CategoryDtypeArg, @@ -394,7 +395,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: ... @@ -430,12 +431,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, @@ -471,7 +472,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 @@ -488,7 +489,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 @@ -522,8 +523,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, ...]: ... diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index dd9f12471..92c60cff9 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/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] diff --git a/pandas-stubs/core/indexes/range.pyi b/pandas-stubs/core/indexes/range.pyi index e96204c36..5ddd3713e 100644 --- a/pandas-stubs/core/indexes/range.pyi +++ b/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]: ... diff --git a/pyproject.toml b/pyproject.toml index c7279eb73..a4d2f546b 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/indexes/test_categoricalindex.py b/tests/indexes/test_categoricalindex.py new file mode 100644 index 000000000..7a9990582 --- /dev/null +++ b/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)