type setitem

Author: MarcoGorelli

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 = ...,

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,