diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index c45da47c1..3cc06ca3c 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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 diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 56d80c24a..2f09f6da8 100644 --- a/pandas-stubs/core/generic.pyi +++ b/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 diff --git a/pandas-stubs/core/indexing.pyi b/pandas-stubs/core/indexing.pyi index 62fc3435a..b3b40b6e1 100644 --- a/pandas-stubs/core/indexing.pyi +++ b/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: ... diff --git a/pyproject.toml b/pyproject.toml index 35fd38a6d..0663fc85d 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/frame/test_indexing.py b/tests/frame/test_indexing.py index 249ddd127..8507c17a3 100644 --- a/tests/frame/test_indexing.py +++ b/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] diff --git a/tests/series/test_indexing.py b/tests/series/test_indexing.py index ddef397ed..da52a46de 100644 --- a/tests/series/test_indexing.py +++ b/tests/series/test_indexing.py @@ -8,6 +8,7 @@ from tests import ( PD_LTE_23, + TYPE_CHECKING_INVALID_USAGE, check, pytest_warns_bounded, ) @@ -31,8 +32,13 @@ 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: @@ -40,8 +46,10 @@ def test_types_loc_at() -> None: 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: