From 0b17616d132d63479d3fdca2dd15cf65d9af9730 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 6 Nov 2025 15:03:01 +0000 Subject: [PATCH 01/10] type asof --- pandas-stubs/core/frame.pyi | 6 +++++- pandas-stubs/core/series.pyi | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index a3d6f70ab..331d60463 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1844,7 +1844,11 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): **kwargs: Any, ) -> Series[_bool]: ... @final - def asof(self, where, subset: _str | list[_str] | None = None) -> Self: ... + def asof( + self, + where: Scalar | AnyArrayLike | Sequence[Scalar], + subset: Hashable | list[Hashable] | None = None, + ) -> Self: ... @final def asfreq( self, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 32b6bc8a5..9b21960c4 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1523,8 +1523,8 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @final def asof( self, - where: Scalar | Sequence[Scalar], - subset: _str | Sequence[_str] | None = None, + where: Scalar | AnyArrayLike | Sequence[Scalar], + subset: None = None, ) -> Scalar | Series[S1]: ... @overload def clip( # pyright: ignore[reportOverlappingOverload] From a05e17f70ee255624e5f07eee751b12676f4eb1b Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 6 Nov 2025 15:19:12 +0000 Subject: [PATCH 02/10] type setitem --- pandas-stubs/core/frame.pyi | 68 +++++++++++++++++++++++++++++++++--- pandas-stubs/core/series.pyi | 4 +-- 2 files changed, 65 insertions(+), 7 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 331d60463..deb20a48b 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -37,6 +37,7 @@ from pandas import ( Timestamp, ) from pandas.core.arraylike import OpsMixin +from pandas.core.base import IndexOpsMixin from pandas.core.generic import NDFrame from pandas.core.groupby.generic import DataFrameGroupBy from pandas.core.indexers import BaseIndexer @@ -90,6 +91,7 @@ from pandas._typing import ( ArrayLike, AstypeArg, Axes, + AxesData, Axis, AxisColumn, AxisIndex, @@ -191,6 +193,8 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]): | tuple[slice] ), ) -> _T: ... + + # Keep in sync with `DataFrame.__setitem__` def __setitem__( self, idx: ( @@ -203,7 +207,7 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]): ), value: ( Scalar - | Series + | IndexOpsMixin | DataFrame | np.ndarray | NAType @@ -267,6 +271,8 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): ) -> Series: ... @overload def __getitem__(self, idx: tuple[Scalar, slice]) -> Series | _T: ... + + # Keep in sync with `DataFrame.__setitem__` @overload def __setitem__( self, @@ -278,7 +284,7 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): | NAType | NaTType | ArrayLike - | Series + | IndexOpsMixin | DataFrame | list | Mapping[Hashable, Scalar | NAType | NaTType] @@ -322,7 +328,7 @@ class _AtIndexerFrame(_AtIndexer): | NAType | NaTType | ArrayLike - | Series + | IndexOpsMixin | DataFrame | list | Mapping[Hashable, Scalar | NAType | NaTType] @@ -794,7 +800,57 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def isetitem( self, loc: int | Sequence[int], value: Scalar | ArrayLike | list[Any] ) -> None: ... - def __setitem__(self, key, value) -> None: ... + + # Keep in sync with `_iLocIndexerFrame.__setitem__` + @overload + def __setitem__( + self, + idx: ( + int + | IndexType + | tuple[int, int] + | tuple[IndexType, int] + | tuple[IndexType, IndexType] + | tuple[int, IndexType] + ), + value: ( + Scalar + | IndexOpsMixin + | DataFrame + | np.ndarray + | NAType + | NaTType + | Mapping[Hashable, Scalar | NAType | NaTType] + | None + ), + ) -> None: ... + # Keep in sync with `_LocIndexerFrame.__setitem__` + @overload + def __setitem__( + self, + idx: ( + MaskType | StrLike | _IndexSliceTuple | list[ScalarT] | IndexingInt | slice + ), + value: ( + Scalar + | NAType + | NaTType + | ArrayLike + | IndexOpsMixin + | DataFrame + | list + | Mapping[Hashable, Scalar | NAType | NaTType] + | None + ), + ) -> None: ... + @overload + def __setitem__( + self, + idx: tuple[_IndexSliceTuple, Hashable], + value: ( + Scalar | NAType | NaTType | ArrayLike | IndexOpsMixin | list | dict | None + ), + ) -> None: ... @overload def query( self, @@ -2398,7 +2454,9 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): **kwargs: Any, ) -> Series: ... # Not actually positional, but used to handle removal of deprecated - def set_axis(self, labels, *, axis: Axis = ..., copy: _bool = ...) -> Self: ... + def set_axis( + self, labels: AxesData, *, axis: Axis = 0, copy: _bool = ... + ) -> Self: ... def skew( self, axis: Axis | None = ..., diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 9b21960c4..8088df68b 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -535,14 +535,14 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @final def __getattr__(self, name: _str) -> S1: ... - # Keep in sync with `iLocIndexerSeries.__getitem__` + # Keep in sync with `_iLocIndexerSeries.__getitem__` @overload def __getitem__(self, idx: IndexingInt) -> S1: ... @overload def __getitem__( self, idx: Index | Series | slice | np_ndarray_anyint ) -> Series[S1]: ... - # Keep in sync with `LocIndexerSeries.__getitem__` + # Keep in sync with `_LocIndexerSeries.__getitem__` @overload def __getitem__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] self, From b66d0c85a56f26b27985c9e4e8dc7651f710536a Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 6 Nov 2025 15:49:53 +0000 Subject: [PATCH 03/10] wip# --- pandas-stubs/core/frame.pyi | 11 ++++++++++- pandas-stubs/core/indexing.pyi | 4 ++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index deb20a48b..c1b9c9e27 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -816,6 +816,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): value: ( Scalar | IndexOpsMixin + | Sequence[Scalar] | DataFrame | np.ndarray | NAType @@ -837,6 +838,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | NaTType | ArrayLike | IndexOpsMixin + | Sequence[Scalar] | DataFrame | list | Mapping[Hashable, Scalar | NAType | NaTType] @@ -848,7 +850,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): self, idx: tuple[_IndexSliceTuple, Hashable], value: ( - Scalar | NAType | NaTType | ArrayLike | IndexOpsMixin | list | dict | None + Scalar + | NAType + | NaTType + | ArrayLike + | IndexOpsMixin + | Sequence[Scalar] + | dict + | None ), ) -> None: ... @overload diff --git a/pandas-stubs/core/indexing.pyi b/pandas-stubs/core/indexing.pyi index 832e2618a..868d10afc 100644 --- a/pandas-stubs/core/indexing.pyi +++ b/pandas-stubs/core/indexing.pyi @@ -3,7 +3,7 @@ from typing import ( TypeVar, ) -from pandas.core.indexes.api import Index +from pandas.core.base import IndexOpsMixin from pandas._libs.indexing import _NDFrameIndexerBase from pandas._typing import ( @@ -13,7 +13,7 @@ from pandas._typing import ( ) _IndexSliceTuple: TypeAlias = tuple[ - Index | MaskType | Scalar | list[ScalarT] | slice | tuple[Scalar, ...], ... + IndexOpsMixin | MaskType | Scalar | list[ScalarT] | slice | tuple[Scalar, ...], ... ] _IndexSliceUnion: TypeAlias = slice | _IndexSliceTuple From 6c91caaee3f7de24455d273662ccd4519a2bc005 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 11 Nov 2025 15:02:00 +0000 Subject: [PATCH 04/10] fixup setitem --- pandas-stubs/core/frame.pyi | 8 ++++++++ tests/test_frame.py | 5 +++++ 2 files changed, 13 insertions(+) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 1bc16cb18..dc39ca8c1 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -866,6 +866,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | None ), ) -> None: ... + # Extra cases not supported by `_LocIndexerFrame.__setitem__` / + # `_iLocIndexerFrame.__setitem__`. + @overload + def __setitem__( + self, + idx: IndexOpsMixin | DataFrame, + value: Scalar | NAType | NaTType | ArrayLike | Series | list | dict | None, + ): ... @overload def query( self, diff --git a/tests/test_frame.py b/tests/test_frame.py index eb2960b0f..abe8db2ad 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -292,6 +292,11 @@ def test_types_setitem() -> None: df[5] = [5, 6] df[["col1", "col2"]] = [[1, 2], [3, 4]] df[s] = [5, 6] + df["col1":"col2"] = [5, 6] + df.loc[:, s] = [5, 6] + df[[0, 1]] = [5, 6] + df["col1"] = [5, 6] + df[df["col1"] > 1] = [5, 6, 7] df[a] = [[1, 2], [3, 4]] df[i] = [8, 9] From 9538b55519b4803d0ac3f2f7dff47ea29f80cde1 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 11 Nov 2025 15:43:15 +0000 Subject: [PATCH 05/10] remove invalid case --- tests/test_frame.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index abe8db2ad..53662dfbf 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -292,7 +292,6 @@ def test_types_setitem() -> None: df[5] = [5, 6] df[["col1", "col2"]] = [[1, 2], [3, 4]] df[s] = [5, 6] - df["col1":"col2"] = [5, 6] df.loc[:, s] = [5, 6] df[[0, 1]] = [5, 6] df["col1"] = [5, 6] From 02f940a797f1f50527a066e93ae67966a802a6c3 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 11 Nov 2025 16:42:04 +0000 Subject: [PATCH 06/10] remove test case which adds extra columns (thus making subsequent tests fail) --- tests/test_frame.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 53662dfbf..63fa9e65e 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -293,7 +293,6 @@ def test_types_setitem() -> None: df[["col1", "col2"]] = [[1, 2], [3, 4]] df[s] = [5, 6] df.loc[:, s] = [5, 6] - df[[0, 1]] = [5, 6] df["col1"] = [5, 6] df[df["col1"] > 1] = [5, 6, 7] df[a] = [[1, 2], [3, 4]] From f2ba8740518f6062ed9de09bd32fdd48ee6c8291 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 11 Nov 2025 18:39:41 +0000 Subject: [PATCH 07/10] use Sequence[Sequence[Scalar]] --- pandas-stubs/core/frame.pyi | 38 +++++++++++++++++++++++++++------- pandas-stubs/core/indexing.pyi | 4 ++-- 2 files changed, 32 insertions(+), 10 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index dc39ca8c1..aa56ece07 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -292,7 +292,8 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): | ArrayLike | IndexOpsMixin | DataFrame - | list + | Sequence[Scalar] + | Sequence[Sequence[Scalar]] | Mapping[Hashable, Scalar | NAType | NaTType] | None ), @@ -301,7 +302,17 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): def __setitem__( self, idx: tuple[_IndexSliceTuple, Hashable], - value: Scalar | NAType | NaTType | ArrayLike | Series | list | dict | None, + value: ( + Scalar + | NAType + | NaTType + | ArrayLike + | Series + | Sequence[Scalar] + | Sequence[Sequence[Scalar]] + | Mapping[Hashable, Scalar | NAType | NaTType] + | None + ), ) -> None: ... class _iAtIndexerFrame(_iAtIndexer): @@ -336,7 +347,8 @@ class _AtIndexerFrame(_AtIndexer): | ArrayLike | IndexOpsMixin | DataFrame - | list + | Sequence[Scalar] + | Sequence[Sequence[Scalar]] | Mapping[Hashable, Scalar | NAType | NaTType] | None ), @@ -824,7 +836,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | IndexOpsMixin | Sequence[Scalar] | DataFrame - | np.ndarray + | np_ndarray | NAType | NaTType | Mapping[Hashable, Scalar | NAType | NaTType] @@ -845,8 +857,8 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | ArrayLike | IndexOpsMixin | Sequence[Scalar] + | Sequence[Sequence[Scalar]] | DataFrame - | list | Mapping[Hashable, Scalar | NAType | NaTType] | None ), @@ -862,7 +874,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | ArrayLike | IndexOpsMixin | Sequence[Scalar] - | dict + | Mapping[Hashable, Scalar | NAType | NaTType] | None ), ) -> None: ... @@ -872,8 +884,18 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def __setitem__( self, idx: IndexOpsMixin | DataFrame, - value: Scalar | NAType | NaTType | ArrayLike | Series | list | dict | None, - ): ... + value: ( + Scalar + | NAType + | NaTType + | ArrayLike + | Series + | Sequence[Scalar] + | Sequence[Sequence[Scalar]] + | Mapping[Hashable, Scalar | NAType | NaTType] + | None + ), + ) -> None: ... @overload def query( self, diff --git a/pandas-stubs/core/indexing.pyi b/pandas-stubs/core/indexing.pyi index 868d10afc..62fc3435a 100644 --- a/pandas-stubs/core/indexing.pyi +++ b/pandas-stubs/core/indexing.pyi @@ -1,3 +1,4 @@ +from collections.abc import Sequence from typing import ( TypeAlias, TypeVar, @@ -9,11 +10,10 @@ from pandas._libs.indexing import _NDFrameIndexerBase from pandas._typing import ( MaskType, Scalar, - ScalarT, ) _IndexSliceTuple: TypeAlias = tuple[ - IndexOpsMixin | MaskType | Scalar | list[ScalarT] | slice | tuple[Scalar, ...], ... + IndexOpsMixin | MaskType | Scalar | Sequence[Scalar] | slice, ... ] _IndexSliceUnion: TypeAlias = slice | _IndexSliceTuple From 57585129385494f0dc4ae5178ef92f9519f70c49 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 12 Nov 2025 10:56:34 +0000 Subject: [PATCH 08/10] re-sync --- pandas-stubs/core/frame.pyi | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index aa56ece07..da456d3cb 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -214,6 +214,7 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]): value: ( Scalar | IndexOpsMixin + | Sequence[Scalar] | DataFrame | np_ndarray | NAType @@ -291,9 +292,9 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): | NaTType | ArrayLike | IndexOpsMixin - | DataFrame | Sequence[Scalar] | Sequence[Sequence[Scalar]] + | DataFrame | Mapping[Hashable, Scalar | NAType | NaTType] | None ), @@ -307,7 +308,7 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): | NAType | NaTType | ArrayLike - | Series + | IndexOpsMixin | Sequence[Scalar] | Sequence[Sequence[Scalar]] | Mapping[Hashable, Scalar | NAType | NaTType] @@ -874,6 +875,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | ArrayLike | IndexOpsMixin | Sequence[Scalar] + | Sequence[Sequence[Scalar]] | Mapping[Hashable, Scalar | NAType | NaTType] | None ), From 6dc4c965b80896ff0258b041bc2243a9a11b4865 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 12 Nov 2025 11:00:38 +0000 Subject: [PATCH 09/10] widen Series to IndexOpsMixin --- pandas-stubs/core/frame.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index da456d3cb..3b9251752 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -891,7 +891,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): | NAType | NaTType | ArrayLike - | Series + | IndexOpsMixin | Sequence[Scalar] | Sequence[Sequence[Scalar]] | Mapping[Hashable, Scalar | NAType | NaTType] From 16a2e683a35f823d0c2a2ae6408eb3c0108f9a30 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 12 Nov 2025 12:15:00 +0000 Subject: [PATCH 10/10] remove deprecated `copy` --- pandas-stubs/core/frame.pyi | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 3b9251752..da31cff03 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -2555,9 +2555,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): **kwargs: Any, ) -> Series: ... # Not actually positional, but used to handle removal of deprecated - def set_axis( - self, labels: AxesData, *, axis: Axis = 0, copy: _bool = ... - ) -> Self: ... + def set_axis(self, labels: AxesData, *, axis: Axis = 0) -> Self: ... def skew( self, axis: Axis | None = ...,