Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 7 additions & 28 deletions pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions pandas-stubs/core/generic.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
54 changes: 26 additions & 28 deletions pandas-stubs/core/indexing.pyi
Original file line number Diff line number Diff line change
@@ -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,
)
Expand All @@ -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: ...
4 changes: 0 additions & 4 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
30 changes: 29 additions & 1 deletion tests/frame/test_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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]
8 changes: 8 additions & 0 deletions tests/series/test_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

from tests import (
PD_LTE_23,
TYPE_CHECKING_INVALID_USAGE,
check,
pytest_warns_bounded,
)
Expand All @@ -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:
Expand Down