type indexing.pyi (#1554)

* type indexing.pyi

* remove indexingmixin

* extra test case

* move test location, add some `TYPE_CHECKING_INVALID_USAGE`

Author: MarcoGorelli

pandas-stubs/core/frame.pyi

@@ -323,42 +323,21 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
     ) -> None: ...
 
 class _iAtIndexerFrame(_iAtIndexer):
-    def __getitem__(self, key: tuple[int, int]) -> Scalar: ...
-    def __setitem__(
+    def __getitem__(self, key: tuple[int, int]) -> Scalar: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+    def __setitem__(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         key: tuple[int, int],
         value: Scalar | NAType | NaTType | None,
     ) -> None: ...
 
 class _AtIndexerFrame(_AtIndexer):
-    def __getitem__(
-        self,
-        key: tuple[
-            int
-            | StrLike
-            | Timestamp
-            | tuple[Scalar, ...]
-            | Callable[[DataFrame], ScalarT],
-            int | StrLike | tuple[Scalar, ...],
-        ],
+    def __getitem__(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+        self, key: tuple[Hashable, Hashable]
     ) -> Scalar: ...
-    def __setitem__(
+    def __setitem__(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
         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/generic.pyi

@@ -22,7 +22,6 @@ from typing import (
 )
 
 from pandas import Index
-import pandas.core.indexing as indexing
 from pandas.core.resample import DatetimeIndexResampler
 from pandas.core.series import Series
 import sqlalchemy.engine
@@ -63,7 +62,7 @@ from pandas._typing import (
 from pandas.io.pytables import HDFStore
 from pandas.io.sql import SQLTable
 
-class NDFrame(indexing.IndexingMixin):
+class NDFrame:
     __hash__: ClassVar[None]  # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride]
 
     @final

pandas-stubs/core/indexing.pyi

@@ -1,13 +1,21 @@
-from collections.abc import Sequence
+from collections.abc import (
+    Hashable,
+    Sequence,
+)
 from typing import (
     TypeAlias,
     TypeVar,
 )
 
 from pandas.core.base import IndexOpsMixin
+from typing_extensions import Self
 
 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,
     MaskType,
     Scalar,
 )
@@ -27,31 +35,21 @@ class _IndexSlice:
 
 IndexSlice: _IndexSlice
 
-class IndexingMixin:
-    @property
-    def iloc(self) -> _iLocIndexer: ...
-    @property
-    def loc(self) -> _LocIndexer: ...
-    @property
-    def at(self) -> _AtIndexer: ...
-    @property
-    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_indexing.py

@@ -26,7 +26,10 @@
 
 from pandas._typing import Scalar
 
-from tests import check
+from tests import (
+    TYPE_CHECKING_INVALID_USAGE,
+    check,
+)
 
 
 def test_types_getitem() -> None:
@@ -592,3 +595,28 @@ def test_frame_ndarray_assignmment() -> None:
 
     df_b = pd.DataFrame({"a": [0.0] * 10, "b": [1.0] * 10})
     df_b.iloc[:, :] = np.array([[-1.0, np.inf]] * 10)
+
+
+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")
+
+    mi = pd.MultiIndex.from_arrays([[2, 3], [4, 5]])
+    df = pd.DataFrame(data={"col1": [1.6, 2], "col2": [3, 4]}, index=mi)
+
+    check(assert_type(df.at[(2, 4), "col1"], Scalar), float)
+    df.at[(2, 4), "col1"] = 999
+    df.at[(2, 4), "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")
+    if TYPE_CHECKING_INVALID_USAGE:
+        df.iat[(0,), 0] = 999  # type: ignore[index]  # pyright: ignore[reportArgumentType]

tests/series/test_indexing.py

@@ -8,6 +8,7 @@
 
 from tests import (
     PD_LTE_23,
+    TYPE_CHECKING_INVALID_USAGE,
     check,
     pytest_warns_bounded,
 )
@@ -31,17 +32,24 @@ 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
+
+    if TYPE_CHECKING_INVALID_USAGE:
+        s.iat[0, 0]  # type: ignore[index]  # pyright: ignore[reportArgumentType]
 
 
 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: