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 1/2] 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 2/2] 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,