From 0c6c60b782632ac2080f8958f8a9bd2f4c2602e8 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Mon, 1 Dec 2025 18:25:14 +0000 Subject: [PATCH 01/12] type Index view and drop --- pandas-stubs/_typing.pyi | 2 ++ pandas-stubs/core/indexes/base.pyi | 14 ++++++++++++-- pyproject.toml | 6 +++++- tests/indexes/test_indexes.py | 26 ++++++++++++++++++++++++++ 4 files changed, 45 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 1350b7c1b..f7dcd1cfc 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -950,6 +950,8 @@ np_1darray_dt: TypeAlias = np_1darray[np.datetime64] np_1darray_td: TypeAlias = np_1darray[np.timedelta64] np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] +NDArrayT = TypeVar("NDArrayT", bound=np.ndarray) + DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) KeysArgType: TypeAlias = Any ListLikeT = TypeVar("ListLikeT", bound=ListLike) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 09bb37a8a..2d7e44a77 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -102,6 +102,7 @@ from pandas._typing import ( Level, MaskType, NaPosition, + NDArrayT, NumpyFloatNot16DtypeArg, PandasAstypeFloatDtypeArg, PandasFloatDtypeArg, @@ -374,7 +375,12 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): def dtype(self) -> DtypeObj: ... @final def ravel(self, order: _str = "C") -> Self: ... - def view(self, cls=...): ... + @overload + def view(self, cls: None = None) -> Self: ... + @overload + def view(self, cls: type[NDArrayT]) -> NDArrayT: ... + @overload + def view(self, cls: Dtype) -> ArrayLike: ... @overload def astype( self, @@ -596,7 +602,11 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): 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: ... + def drop( + self, + labels: IndexOpsMixin | np_1darray | Iterable[Hashable], + errors: IgnoreRaise = "raise", + ) -> Self: ... @property def shape(self) -> tuple[int, ...]: ... # Extra methods from old stubs diff --git a/pyproject.toml b/pyproject.toml index 46fc93532..ec6c53470 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -237,7 +237,11 @@ ignore = [ # TODO: remove when _libs is fully typed "ANN001", "ANN201", "ANN204", "ANN206", ] -"*base.pyi" = [ +"*core/base.pyi" = [ + # TODO: remove when base.pyi's are fully typed + "ANN001", "ANN201", "ANN204", "ANN206", +] +"*excel/_base.pyi" = [ # TODO: remove when base.pyi's are fully typed "ANN001", "ANN201", "ANN204", "ANN206", ] diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index eda08354e..a82f646e7 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -22,6 +22,7 @@ assert_type, ) +from pandas._typing import ArrayLike # noqa: F401 from pandas._typing import Dtype # noqa: F401 from pandas._typing import Scalar # noqa: F401 @@ -1667,3 +1668,28 @@ def test_index_slice_locs() -> None: start, end = idx.slice_locs(0, 1) check(assert_type(start, np.intp | int), np.integer) check(assert_type(end, np.intp | int), int) + + +def test_index_view() -> None: + ind = pd.Index([1, 2]) + check(assert_type(ind.view("int64"), ArrayLike), np_1darray_int64) + check(assert_type(ind.view(), "pd.Index[int]"), pd.Index) + check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) + + class MyArray(np.ndarray): ... + + check(assert_type(ind.view(MyArray), MyArray), MyArray) + + +def test_index_drop() -> None: + ind = pd.Index([1, 2, 3]) + check(assert_type(ind.drop([1, 2]), "pd.Index[int]"), pd.Index, np.integer) + check( + assert_type(ind.drop(pd.Index([1, 2])), "pd.Index[int]"), pd.Index, np.integer + ) + check( + assert_type(ind.drop(pd.Series([1, 2])), "pd.Index[int]"), pd.Index, np.integer + ) + check( + assert_type(ind.drop(np.array([1, 2])), "pd.Index[int]"), pd.Index, np.integer + ) From cff1f764c0bd3b2ff0cf76e94a2b1c92611fc635 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Mon, 1 Dec 2025 21:07:49 +0000 Subject: [PATCH 02/12] note mypy vs pyright discrepancy --- pandas-stubs/_typing.pyi | 2 +- tests/indexes/test_indexes.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index f7dcd1cfc..c37609123 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -950,7 +950,7 @@ np_1darray_dt: TypeAlias = np_1darray[np.datetime64] np_1darray_td: TypeAlias = np_1darray[np.timedelta64] np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] -NDArrayT = TypeVar("NDArrayT", bound=np.ndarray) +NDArrayT = TypeVar("NDArrayT", bound=np_ndarray) DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) KeysArgType: TypeAlias = Any diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index a82f646e7..40ae63068 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -37,6 +37,7 @@ np_1darray_bool, np_1darray_int64, np_1darray_intp, + np_ndarray, np_ndarray_dt, ) @@ -1674,7 +1675,10 @@ def test_index_view() -> None: ind = pd.Index([1, 2]) check(assert_type(ind.view("int64"), ArrayLike), np_1darray_int64) check(assert_type(ind.view(), "pd.Index[int]"), pd.Index) - check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) + # mypy and pyright differ here in what they report: + # - mypy: ndarray[Any, Any]" + # - pyright: ndarray[tuple[Any, ...], dtype[Any]] + check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] class MyArray(np.ndarray): ... From 6e67867a4fe609634f229ececf948f04e4efaa26 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Mon, 1 Dec 2025 21:24:27 +0000 Subject: [PATCH 03/12] old numpy --- tests/indexes/test_indexes.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 40ae63068..5efd97a69 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -2,6 +2,7 @@ from collections.abc import Hashable import datetime as dt +import sys from typing import ( Any, cast, @@ -1675,10 +1676,13 @@ def test_index_view() -> None: ind = pd.Index([1, 2]) check(assert_type(ind.view("int64"), ArrayLike), np_1darray_int64) check(assert_type(ind.view(), "pd.Index[int]"), pd.Index) - # mypy and pyright differ here in what they report: - # - mypy: ndarray[Any, Any]" - # - pyright: ndarray[tuple[Any, ...], dtype[Any]] - check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] + # on NumPy<1.23, we get: + # error: "assert_type" mismatch: expected "ndarray[tuple[Any, ...], dtype[Any]]" but received "ndarray[Unknown, Unknown]" (reportAssertTypeFailure) + if sys.version_info >= (3, 11): + # mypy and pyright differ here in what they report: + # - mypy: ndarray[Any, Any]" + # - pyright: ndarray[tuple[Any, ...], dtype[Any]] + check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] class MyArray(np.ndarray): ... From 24c87d6649fabf20dcf4e12678dfdc6b0f0799c9 Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 2 Dec 2025 09:53:23 +0000 Subject: [PATCH 04/12] more precise comments, laxer arg types, return 1d np array Co-authored-by: Yi-Fan Wang --- pandas-stubs/_typing.pyi | 2 +- pandas-stubs/core/indexes/base.pyi | 4 ++-- pyproject.toml | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index c37609123..f7dcd1cfc 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -950,7 +950,7 @@ np_1darray_dt: TypeAlias = np_1darray[np.datetime64] np_1darray_td: TypeAlias = np_1darray[np.timedelta64] np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] -NDArrayT = TypeVar("NDArrayT", bound=np_ndarray) +NDArrayT = TypeVar("NDArrayT", bound=np.ndarray) DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) KeysArgType: TypeAlias = Any diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2d7e44a77..2ccde53f9 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -380,7 +380,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @overload def view(self, cls: type[NDArrayT]) -> NDArrayT: ... @overload - def view(self, cls: Dtype) -> ArrayLike: ... + def view(self, cls: NumpyNotTimeDtypeArg | NumpyTimedeltaDtypeArg | NumpyTimestampDtypeArg) -> np_1darray: ... @overload def astype( self, @@ -604,7 +604,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): def insert(self, loc: int, item: object) -> Index: ... def drop( self, - labels: IndexOpsMixin | np_1darray | Iterable[Hashable], + labels: IndexOpsMixin | np_ndarray | Iterable[Hashable], errors: IgnoreRaise = "raise", ) -> Self: ... @property diff --git a/pyproject.toml b/pyproject.toml index ec6c53470..437f0de72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -238,11 +238,11 @@ ignore = [ "ANN001", "ANN201", "ANN204", "ANN206", ] "*core/base.pyi" = [ - # TODO: remove when base.pyi's are fully typed + # TODO: remove when core/base.pyi is fully typed "ANN001", "ANN201", "ANN204", "ANN206", ] "*excel/_base.pyi" = [ - # TODO: remove when base.pyi's are fully typed + # TODO: remove when excel/_base.pyi is fully typed "ANN001", "ANN201", "ANN204", "ANN206", ] "scripts/*" = [ From 348663f6d6190bbc3e9b3e3937f5a65a21fc306e Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 2 Dec 2025 10:15:47 +0000 Subject: [PATCH 05/12] fixup --- pandas-stubs/core/indexes/base.pyi | 8 +++++++- tests/indexes/test_indexes.py | 4 +++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2ccde53f9..6c6c10aa7 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -104,6 +104,9 @@ from pandas._typing import ( NaPosition, NDArrayT, NumpyFloatNot16DtypeArg, + NumpyNotTimeDtypeArg, + NumpyTimedeltaDtypeArg, + NumpyTimestampDtypeArg, PandasAstypeFloatDtypeArg, PandasFloatDtypeArg, PyArrowFloatDtypeArg, @@ -380,7 +383,10 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @overload def view(self, cls: type[NDArrayT]) -> NDArrayT: ... @overload - def view(self, cls: NumpyNotTimeDtypeArg | NumpyTimedeltaDtypeArg | NumpyTimestampDtypeArg) -> np_1darray: ... + def view( + self, + cls: NumpyNotTimeDtypeArg | NumpyTimedeltaDtypeArg | NumpyTimestampDtypeArg, + ) -> np_1darray: ... @overload def astype( self, diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 5efd97a69..6bb233662 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1674,7 +1674,7 @@ def test_index_slice_locs() -> None: def test_index_view() -> None: ind = pd.Index([1, 2]) - check(assert_type(ind.view("int64"), ArrayLike), np_1darray_int64) + check(assert_type(ind.view("int64"), np_1darray), np_1darray) check(assert_type(ind.view(), "pd.Index[int]"), pd.Index) # on NumPy<1.23, we get: # error: "assert_type" mismatch: expected "ndarray[tuple[Any, ...], dtype[Any]]" but received "ndarray[Unknown, Unknown]" (reportAssertTypeFailure) @@ -1683,6 +1683,8 @@ def test_index_view() -> None: # - mypy: ndarray[Any, Any]" # - pyright: ndarray[tuple[Any, ...], dtype[Any]] check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] + else: + check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] class MyArray(np.ndarray): ... From 05af6b0943d7f6c2007eab2223d12cc897fc29d8 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 2 Dec 2025 10:25:01 +0000 Subject: [PATCH 06/12] fixup --- tests/indexes/test_indexes.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 6bb233662..1510731b2 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -23,7 +23,6 @@ assert_type, ) -from pandas._typing import ArrayLike # noqa: F401 from pandas._typing import Dtype # noqa: F401 from pandas._typing import Scalar # noqa: F401 @@ -1684,7 +1683,7 @@ def test_index_view() -> None: # - pyright: ndarray[tuple[Any, ...], dtype[Any]] check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] else: - check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] + check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] # pyright: ignore[reportAssertTypeFailure] class MyArray(np.ndarray): ... From 7d9d43303a95c6d02ba4beb34385e87fe5c6124f Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 2 Dec 2025 10:25:18 +0000 Subject: [PATCH 07/12] Update tests/indexes/test_indexes.py Co-authored-by: Yi-Fan Wang --- tests/indexes/test_indexes.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 6bb233662..939910523 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1703,3 +1703,7 @@ def test_index_drop() -> None: check( assert_type(ind.drop(np.array([1, 2])), "pd.Index[int]"), pd.Index, np.integer ) + check( + assert_type(ind.drop(iter([1, 2])), "pd.Index[int]"), pd.Index, np.integer +Comment view + ) From efd03b331da17f9970f1d503447c0ebba34d60e7 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 2 Dec 2025 10:28:22 +0000 Subject: [PATCH 08/12] fixup syntax --- tests/indexes/test_indexes.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 581acd718..8c248bafc 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1702,7 +1702,4 @@ def test_index_drop() -> None: check( assert_type(ind.drop(np.array([1, 2])), "pd.Index[int]"), pd.Index, np.integer ) - check( - assert_type(ind.drop(iter([1, 2])), "pd.Index[int]"), pd.Index, np.integer -Comment view - ) + check(assert_type(ind.drop(iter([1, 2])), "pd.Index[int]"), pd.Index, np.integer) From e977a5140416e149aca832bc83ce0a1ec8a518ad Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 3 Dec 2025 10:47:44 +0000 Subject: [PATCH 09/12] use np.ndarray in test --- tests/indexes/test_indexes.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 8c248bafc..bf1c10a8c 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -37,7 +37,6 @@ np_1darray_bool, np_1darray_int64, np_1darray_intp, - np_ndarray, np_ndarray_dt, ) @@ -1681,9 +1680,9 @@ def test_index_view() -> None: # mypy and pyright differ here in what they report: # - mypy: ndarray[Any, Any]" # - pyright: ndarray[tuple[Any, ...], dtype[Any]] - check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] + check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) # type: ignore[assert-type] else: - check(assert_type(ind.view(np.ndarray), np_ndarray), np.ndarray) # type: ignore[assert-type] # pyright: ignore[reportAssertTypeFailure] + check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) # type: ignore[assert-type] # pyright: ignore[reportAssertTypeFailure] class MyArray(np.ndarray): ... From 16e084079e77dd8d3280d6a05f142571573839de Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 3 Dec 2025 11:05:37 +0000 Subject: [PATCH 10/12] use np.ndarray in test --- tests/indexes/test_indexes.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index bf1c10a8c..00d3957e7 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1682,7 +1682,9 @@ def test_index_view() -> None: # - pyright: ndarray[tuple[Any, ...], dtype[Any]] check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) # type: ignore[assert-type] else: - check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) # type: ignore[assert-type] # pyright: ignore[reportAssertTypeFailure] + check( + assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray + ) # pyright: ignore[reportAssertTypeFailure] class MyArray(np.ndarray): ... From c31d38287877fb071e95830c5aaa4738937bc355 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 3 Dec 2025 11:20:47 +0000 Subject: [PATCH 11/12] use np.ndarray in test --- tests/indexes/test_indexes.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 00d3957e7..4ebf8235e 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1682,9 +1682,7 @@ def test_index_view() -> None: # - pyright: ndarray[tuple[Any, ...], dtype[Any]] check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) # type: ignore[assert-type] else: - check( - assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray - ) # pyright: ignore[reportAssertTypeFailure] + check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray) class MyArray(np.ndarray): ... From 4f16908d4d9f004edeac187bebffcbf30eb4339b Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 3 Dec 2025 11:21:12 +0000 Subject: [PATCH 12/12] use np.ndarray in test --- tests/indexes/test_indexes.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 4ebf8235e..56e1e87ae 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1674,8 +1674,6 @@ def test_index_view() -> None: ind = pd.Index([1, 2]) check(assert_type(ind.view("int64"), np_1darray), np_1darray) check(assert_type(ind.view(), "pd.Index[int]"), pd.Index) - # on NumPy<1.23, we get: - # error: "assert_type" mismatch: expected "ndarray[tuple[Any, ...], dtype[Any]]" but received "ndarray[Unknown, Unknown]" (reportAssertTypeFailure) if sys.version_info >= (3, 11): # mypy and pyright differ here in what they report: # - mypy: ndarray[Any, Any]"