type indexing.pyi

Author: MarcoGorelli

pandas-stubs/core/frame.pyi

@@ -333,32 +333,12 @@ class _iAtIndexerFrame(_iAtIndexer):
 class _AtIndexerFrame(_AtIndexer):
     def __getitem__(
         self,
-        key: tuple[
-            int
-            | StrLike
-            | Timestamp
-            | tuple[Scalar, ...]
-            | Callable[[DataFrame], ScalarT],
-            int | StrLike | tuple[Scalar, ...],
-        ],
+        key: tuple[Hashable, Hashable],
     ) -> Scalar: ...
     def __setitem__(
         self,
-        key: (
-            MaskType | StrLike | _IndexSliceTuple | list[ScalarT] | IndexingInt | slice
-        ),
-        value: (
-            Scalar
-            | NAType
-            | NaTType
-            | ArrayLike
-            | IndexOpsMixin
-            | DataFrame
-            | Sequence[Scalar]
-            | Sequence[Sequence[Scalar]]
-            | Mapping[Hashable, Scalar | NAType | NaTType]
-            | None
-        ),
+        key: tuple[Hashable, Hashable],
+        value: Scalar | NAType | NaTType | None,
     ) -> None: ...
 
 # With mypy 1.14.1 and python 3.12, the second overload needs a type-ignore statement

pandas-stubs/core/indexing.pyi

@@ -1,13 +1,19 @@
 from collections.abc import Sequence
 from typing import (
+    Self,
     TypeAlias,
     TypeVar,
 )
 
 from pandas.core.base import IndexOpsMixin
 
 from pandas._libs.indexing import _NDFrameIndexerBase
+from pandas._libs.missing import NAType
+from pandas._libs.tslibs.nattype import NaTType
 from pandas._typing import (
+    Axis,
+    AxisInt,
+    Hashable,
     MaskType,
     Scalar,
 )
@@ -38,20 +44,20 @@ class IndexingMixin:
     def iat(self) -> _iAtIndexer: ...
 
 class _NDFrameIndexer(_NDFrameIndexerBase):
-    axis = ...
-    def __call__(self, axis=...): ...
-    def __getitem__(self, key): ...
-    def __setitem__(self, key, value) -> None: ...
-
-class _LocationIndexer(_NDFrameIndexer):
-    def __getitem__(self, key): ...
-
-class _LocIndexer(_LocationIndexer): ...
-class _iLocIndexer(_LocationIndexer): ...
-
-class _ScalarAccessIndexer(_NDFrameIndexerBase):
-    def __getitem__(self, key): ...
-    def __setitem__(self, key, value) -> None: ...
-
-class _AtIndexer(_ScalarAccessIndexer): ...
-class _iAtIndexer(_ScalarAccessIndexer): ...
+    axis: AxisInt | None = None
+    def __call__(self, axis: Axis | None = None) -> Self: ...
+
+class _LocIndexer(_NDFrameIndexer): ...
+class _iLocIndexer(_NDFrameIndexer): ...
+
+class _AtIndexer(_NDFrameIndexerBase):
+    def __getitem__(self, key: Hashable) -> Scalar: ...
+    def __setitem__(
+        self, key: Hashable, value: Scalar | NAType | NaTType | None
+    ) -> None: ...
+
+class _iAtIndexer(_NDFrameIndexerBase):
+    def __getitem__(self, key: int) -> Scalar: ...
+    def __setitem__(
+        self, key: int, value: Scalar | NAType | NaTType | None
+    ) -> None: ...

pyproject.toml

@@ -205,10 +205,6 @@ ignore = [
   "PYI042",  # https://docs.astral.sh/ruff/rules/snake-case-type-alias/
   "ERA001", "PLR0402", "PLC0105"
 ]
-"indexing.pyi" = [
-  # TODO: remove when indexing.pyi is fully typed
-  "ANN001", "ANN201", "ANN204", "ANN206",
-]
 "*computation*" = [
   # TODO: remove when computations are fully typed
   "ANN001", "ANN201", "ANN204", "ANN206",

tests/frame/test_frame.py

@@ -3959,3 +3959,19 @@ def test_frame_index_setter() -> None:
     check(assert_type(df.index, pd.Index), pd.Index)
     df.index = [2, 3]
     check(assert_type(df.index, pd.Index), pd.Index)
+
+
+def test_frame_at() -> None:
+    df = pd.DataFrame(data={"col1": [1.6, 2], "col2": [3, 4]})
+
+    check(assert_type(df.at[0, "col1"], Scalar), float)
+    df.at[0, "col1"] = 999
+    df.at[0, "col1"] = float("nan")
+
+
+def test_frame_iat() -> None:
+    df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
+
+    check(assert_type(df.iat[0, 0], Scalar), np.integer)
+    df.iat[0, 0] = 999
+    df.iat[0, 0] = float("nan")

tests/series/test_indexing.py

@@ -31,17 +31,21 @@ def test_types_iloc_iat() -> None:
     s2 = pd.Series(data=[1, 2])
     s.loc["row1"]
     s.iat[0]
+    s.iat[0] = 999
     s2.loc[0]
     s2.iat[0]
+    s2.iat[0] = None
 
 
 def test_types_loc_at() -> None:
     s = pd.Series(data={"row1": 1, "row2": 2})
     s2 = pd.Series(data=[1, 2])
     s.loc["row1"]
     s.at["row1"]
+    s.at["row1"] = 9
     s2.loc[1]
     s2.at[1]
+    s2.at[1] = 99
 
 
 def test_types_getitem() -> None: