From 6518ad3a8d0f69aa8c24e4de537ba02218a6fe08 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 11:46:43 +0000 Subject: [PATCH 1/8] wip --- pandas-stubs/core/arrays/categorical.pyi | 2 -- pandas-stubs/core/indexes/base.pyi | 11 ++++++----- pandas-stubs/core/indexes/category.pyi | 24 ++++++++---------------- pandas-stubs/core/indexes/range.pyi | 1 - 4 files changed, 14 insertions(+), 24 deletions(-) diff --git a/pandas-stubs/core/arrays/categorical.pyi b/pandas-stubs/core/arrays/categorical.pyi index 667d46819..865624503 100644 --- a/pandas-stubs/core/arrays/categorical.pyi +++ b/pandas-stubs/core/arrays/categorical.pyi @@ -110,7 +110,6 @@ 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 @@ -121,7 +120,6 @@ 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=...): ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index c4bf83138..407ab2398 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: ... @@ -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,8 @@ 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: int | AnyArrayLikeInt | Sequence[int]) -> Self: ... + def insert(self, loc: int, item: S1) -> Self: ... 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..c54e5e46b 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,10 @@ 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, +) class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): codes: np.ndarray = ... @@ -22,27 +24,17 @@ class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): cls, data: Iterable[S1] = ..., categories=..., - ordered=..., - dtype=..., - copy: bool = ..., - name: Hashable = ..., + 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): ... 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]: ... From cb0709fbf405be3c88fba5a96d8040e96b578e3c Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 12:10:40 +0000 Subject: [PATCH 2/8] bit more --- pandas-stubs/core/indexes/base.pyi | 12 ++++++------ pandas-stubs/core/indexes/category.pyi | 1 - 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 407ab2398..2e05361bf 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -431,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, diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index c54e5e46b..3d64efd9a 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/pandas-stubs/core/indexes/category.pyi @@ -37,4 +37,3 @@ class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): def is_monotonic_increasing(self) -> bool: ... @property def is_monotonic_decreasing(self) -> bool: ... - def reindex(self, target, method=..., level=..., limit=..., tolerance=...): ... From 6dc0ea9618351fa368ef38b4d187fb38c27d8f17 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 12:33:45 +0000 Subject: [PATCH 3/8] testing --- tests/indexes/test_categoricalindex.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 tests/indexes/test_categoricalindex.py diff --git a/tests/indexes/test_categoricalindex.py b/tests/indexes/test_categoricalindex.py new file mode 100644 index 000000000..942fc8efa --- /dev/null +++ b/tests/indexes/test_categoricalindex.py @@ -0,0 +1,27 @@ +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"]) + 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) From 0015e537331430f02c6b7ccc939ac6cc42ea5c51 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 14:22:06 +0000 Subject: [PATCH 4/8] more tests --- pandas-stubs/core/indexes/base.pyi | 2 +- tests/indexes/test_categoricalindex.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2e05361bf..bce81746f 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -524,7 +524,7 @@ 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: ... - def insert(self, loc: int, item: S1) -> Self: ... + 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/tests/indexes/test_categoricalindex.py b/tests/indexes/test_categoricalindex.py index 942fc8efa..4b59d134e 100644 --- a/tests/indexes/test_categoricalindex.py +++ b/tests/indexes/test_categoricalindex.py @@ -25,3 +25,16 @@ def test_categoricalindex_reindex() -> None: 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) + + +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) From 5365bfc232e4305daddb9acf20ec7d59d9a9a7ab Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 14:43:30 +0000 Subject: [PATCH 5/8] fixup `insert` --- pandas-stubs/core/arrays/categorical.pyi | 2 ++ pandas-stubs/core/indexes/base.pyi | 3 +++ pandas-stubs/core/indexes/category.pyi | 3 +++ tests/indexes/test_categoricalindex.py | 7 ++++--- 4 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/core/arrays/categorical.pyi b/pandas-stubs/core/arrays/categorical.pyi index 865624503..667d46819 100644 --- a/pandas-stubs/core/arrays/categorical.pyi +++ b/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=...): ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index bce81746f..410fad383 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/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 diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index 3d64efd9a..fef5a052b 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/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] diff --git a/tests/indexes/test_categoricalindex.py b/tests/indexes/test_categoricalindex.py index 4b59d134e..7a9990582 100644 --- a/tests/indexes/test_categoricalindex.py +++ b/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: From 63d54f2819016390e681cc9efdb1c9013a286c79 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 16:52:39 +0000 Subject: [PATCH 6/8] include `np.integer`, as per upstream --- pandas-stubs/core/indexes/base.pyi | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 410fad383..7f9f80dc4 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -523,7 +523,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): def slice_locs( self, start: SliceType = None, end: SliceType = None, step: int | None = None ): ... - def delete(self, loc: int | AnyArrayLikeInt | Sequence[int]) -> Self: ... + def delete( + self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int] + ) -> Self: ... @overload def insert(self, loc: int, item: S1) -> Self: ... @overload From 651cf09fb92b68021ca5b56dfaec772572ab87c3 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 30 Oct 2025 16:55:03 +0000 Subject: [PATCH 7/8] type `categories` too --- pandas-stubs/core/indexes/category.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index fef5a052b..92c60cff9 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/pandas-stubs/core/indexes/category.pyi @@ -13,6 +13,7 @@ from typing_extensions import Self from pandas._typing import ( S1, Dtype, + ListLike, ) class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): @@ -23,7 +24,7 @@ class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): def __new__( cls, data: Iterable[S1] = ..., - categories=..., + categories: ListLike | None = None, ordered: bool | None = None, dtype: Dtype | None = None, copy: bool = False, From 97b4261dd6bd64d9a8d92ab72fdcfdd11f91177c Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Fri, 31 Oct 2025 10:06:43 +0000 Subject: [PATCH 8/8] update pyproject.toml --- pyproject.toml | 4 ---- 1 file changed, 4 deletions(-) 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",