Skip to content

TYP: DataFrame.__setitem__ with None #1550

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
loicdiridollou merged 6 commits into from
Dec 18, 2025
Merged

TYP: DataFrame.__setitem__ with None #1550

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
746c760
setitem with none
cmp0xff Dec 11, 2025
7308eb9
more tests and typings
cmp0xff Dec 15, 2025
fe2890b
Merge branch 'main' of github.com:pandas-dev/pandas-stubs into featur…
cmp0xff Dec 16, 2025
4682c1a
python/mypy#20420 https://github.com/pandas-dev/pandas-stubs/pull/155…
cmp0xff Dec 16, 2025
3f139e5
Merge branch 'main' of github.com:pandas-dev/pandas-stubs into featur…
cmp0xff Dec 16, 2025
9bfe743
Merge branch 'main' of github.com:pandas-dev/pandas-stubs into featur…
cmp0xff Dec 18, 2025
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
2 changes: 2 additions & 0 deletions pandas-stubs/_typing.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ from pandas._libs.tslibs import (
Timedelta,
Timestamp,
)
from pandas._libs.tslibs.nattype import NaTType

from pandas.core.dtypes.dtypes import (
CategoricalDtype,
Expand Down Expand Up @@ -134,6 +135,7 @@ _IndexIterScalar: TypeAlias = (
Scalar: TypeAlias = (
_IndexIterScalar | complex | np.integer | np.floating | np.complexfloating
)
ScalarOrNA: TypeAlias = Scalar | NAType | NaTType | None
IntStrT = TypeVar("IntStrT", int, str)

# timestamp and timedelta convertible types
Expand Down
167 changes: 39 additions & 128 deletions pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,6 @@ import xarray as xr
from pandas._libs.lib import _NoDefaultDoNotUse
from pandas._libs.missing import NAType
from pandas._libs.tslibs import BaseOffset
from pandas._libs.tslibs.nattype import NaTType
from pandas._typing import (
S2,
AggFuncTypeBase,
Expand Down Expand Up @@ -146,6 +145,7 @@ from pandas._typing import (
Renamer,
ReplaceValue,
Scalar,
ScalarOrNA,
ScalarT,
SequenceNotStr,
SeriesByT,
Expand Down Expand Up @@ -181,6 +181,26 @@ _T_MUTABLE_MAPPING_co = TypeVar(
"_T_MUTABLE_MAPPING_co", bound=MutableMapping, covariant=True
)

_iLocSetItemKey: TypeAlias = (
int
| IndexType
| tuple[int, int]
| tuple[IndexType, int]
| tuple[IndexType, IndexType]
| tuple[int, IndexType]
)
_LocSetItemKey: TypeAlias = (
MaskType | Hashable | _IndexSliceTuple | Iterable[Scalar] | IndexingInt | slice
)
_SetItemValueNotDataFrame: TypeAlias = (
ScalarOrNA
| Sequence[ScalarOrNA]
| Sequence[Sequence[ScalarOrNA]]
| Mapping[Any, ScalarOrNA]
| ArrayLike
| IndexOpsMixin
)

class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
@overload
def __getitem__(self, key: tuple[int, int]) -> Scalar: ...
Expand All @@ -202,27 +222,13 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
) -> _T: ...

# Keep in sync with `DataFrame.__setitem__`
@overload
def __setitem__(
self,
key: (
int
| IndexType
| tuple[int, int]
| tuple[IndexType, int]
| tuple[IndexType, IndexType]
| tuple[int, IndexType]
),
value: (
Scalar
| IndexOpsMixin
| Sequence[Scalar]
| DataFrame
| np_ndarray
| NAType
| NaTType
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, key: tuple[slice, Hashable], value: _SetItemValueNotDataFrame
) -> None: ...
@overload
def __setitem__(
self, key: _iLocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
) -> None: ...

class _LocIndexerFrame(_LocIndexer, Generic[_T]):
Expand Down Expand Up @@ -283,61 +289,25 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
# Keep in sync with `DataFrame.__setitem__`
@overload
def __setitem__(
self,
key: tuple[_IndexSliceTuple, Hashable],
value: (
Scalar
| NAType
| NaTType
| ArrayLike
| IndexOpsMixin
| Sequence[Scalar]
| Sequence[Sequence[Scalar]]
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, key: tuple[_IndexSliceTuple, Hashable], value: _SetItemValueNotDataFrame
) -> None: ...
@overload
def __setitem__(
self,
key: (
MaskType
| Hashable
| _IndexSliceTuple
| Iterable[Scalar]
| IndexingInt
| slice
),
value: (
Scalar
| NAType
| NaTType
| ArrayLike
| IndexOpsMixin
| Sequence[Scalar]
| Sequence[Sequence[Scalar]]
| DataFrame
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, key: _LocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
) -> None: ...

class _iAtIndexerFrame(_iAtIndexer):
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,
self, key: tuple[int, int], value: ScalarOrNA
) -> None: ...

class _AtIndexerFrame(_AtIndexer):
def __getitem__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
self, key: tuple[Hashable, Hashable]
) -> Scalar: ...
def __setitem__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
self,
key: tuple[Hashable, Hashable],
value: Scalar | NAType | NaTType | None,
self, key: tuple[Hashable, Hashable], value: ScalarOrNA
) -> None: ...

class _GetItemHack:
Expand Down Expand Up @@ -816,85 +786,26 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
# 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
| Sequence[Scalar]
| DataFrame
| np_ndarray
| NAType
| NaTType
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, idex: tuple[slice, Hashable], value: _SetItemValueNotDataFrame
) -> None: ...
@overload
def __setitem__(
self, idx: _iLocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
) -> None: ...
# Keep in sync with `_LocIndexerFrame.__setitem__`
@overload
def __setitem__(
self,
idx: tuple[_IndexSliceTuple, Hashable],
value: (
Scalar
| NAType
| NaTType
| ArrayLike
| IndexOpsMixin
| Sequence[Scalar]
| Sequence[Sequence[Scalar]]
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, idx: tuple[_IndexSliceTuple, Hashable], value: _SetItemValueNotDataFrame
) -> None: ...
@overload
def __setitem__(
self,
idx: (
MaskType
| Hashable
| _IndexSliceTuple
| Iterable[Scalar]
| IndexingInt
| slice
),
value: (
Scalar
| NAType
| NaTType
| ArrayLike
| IndexOpsMixin
| Sequence[Scalar]
| Sequence[Sequence[Scalar]]
| DataFrame
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, idx: _LocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
) -> None: ...
# Extra cases not supported by `_LocIndexerFrame.__setitem__` /
# `_iLocIndexerFrame.__setitem__`.
@overload
def __setitem__(
self,
idx: IndexOpsMixin | DataFrame,
value: (
Scalar
| NAType
| NaTType
| ArrayLike
| IndexOpsMixin
| Sequence[Scalar]
| Sequence[Sequence[Scalar]]
| Mapping[Hashable, Scalar | NAType | NaTType]
| None
),
self, idx: IndexOpsMixin | DataFrame, value: _SetItemValueNotDataFrame
) -> None: ...
@overload
def query(
Expand Down
39 changes: 28 additions & 11 deletions tests/frame/test_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,11 @@ def test_types_setitem() -> None:
df[a] = [[1, 2], [3, 4]]
df[i] = [8, 9]

df["col1"] = [None, pd.NaT]
# TODO: mypy bug, remove after python/mypy#20420 has been resolved
df[["col1"]] = [[None], [pd.NA]] # type: ignore[assignment,list-item]
df[iter(["col1"])] = [[None], [pd.NA]] # type: ignore[assignment]
Comment on lines +98 to +99
Copy link
Contributor Author

@cmp0xff cmp0xff Dec 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not too sure why they don't work for mypy

Copy link
Contributor Author

@cmp0xff cmp0xff Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4682c1a



def test_types_setitem_mask() -> None:
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4], 5: [6, 7]})
Expand Down Expand Up @@ -368,17 +373,6 @@ def test_isetframe() -> None:
check(assert_type(frame.isetitem([0], [10, 12]), None), type(None))


def test_setitem_none() -> None:
df = pd.DataFrame(
{"A": [1, 2, 3], "B": ["abc", "def", "ghi"]}, index=["x", "y", "z"]
)
df.loc["x", "B"] = None
df.iloc[2, 0] = None
sb = pd.Series([1, 2, 3], dtype=int)
sb.loc["y"] = None
sb.iloc[0] = None


def test_getsetitem_multiindex() -> None:
# GH 466
rows = pd.Index(["project A", "project B", "project C"])
Expand Down Expand Up @@ -418,12 +412,32 @@ def test_frame_setitem_na() -> None:

df.loc[ind, :] = pd.NA
df.iloc[[0, 2], :] = pd.NA
df.at["a", "x"] = pd.NA
df.iat[0, 0] = pd.NA

# reveal_type(df["y"]) gives Series[Any], so we have to cast to tell the
# type checker what kind of type it is when adding to a Timedelta
df["x"] = cast("pd.Series[pd.Timestamp]", df["y"]) + pd.Timedelta(days=3)
df.loc[ind, :] = pd.NaT
df.iloc[[0, 2], :] = pd.NaT
df.at["a", "y"] = pd.NaT
df.iat[0, 0] = pd.NaT

df.loc["a", "x"] = None
df.iloc[2, 0] = None
df.at["a", "y"] = None
df.iat[0, 0] = None

df.loc[:, "x"] = [None, pd.NA, pd.NaT]
df.iloc[:, 0] = [None, pd.NA, pd.NaT]

# TODO: mypy bug, remove after python/mypy#20420 has been resolved
df.loc[:, ["x"]] = [[None], [pd.NA], [pd.NaT]] # type: ignore[assignment,index]
df.iloc[:, [0]] = [[None], [pd.NA], [pd.NaT]] # type: ignore[assignment,index]

# TODO: mypy bug, remove after python/mypy#20420 has been resolved
df.loc[:, iter(["x"])] = [[None], [pd.NA], [pd.NaT]] # type: ignore[assignment,index]
df.iloc[:, iter([0])] = [[None], [pd.NA], [pd.NaT]] # type: ignore[assignment,index]


def test_loc_set() -> None:
Expand Down Expand Up @@ -574,6 +588,9 @@ def test_df_loc_dict() -> None:
df.iloc[0] = {"X": 0}
check(assert_type(df, pd.DataFrame), pd.DataFrame)

df.loc[0] = {None: None, pd.NA: pd.NA, pd.NaT: pd.NaT}
df.iloc[0] = {None: None, pd.NA: pd.NA, pd.NaT: pd.NaT}


def test_iloc_npint() -> None:
# GH 69
Expand Down
4 changes: 4 additions & 0 deletions tests/series/test_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,10 @@ def test_series_setitem_na() -> None:
s2.loc[ind] = pd.NaT
s2.iloc[[0, 2]] = pd.NaT

sb = pd.Series([1, 2, 3], dtype=int)
sb.loc["y"] = None
sb.iloc[0] = None


def test_slice_timestamp() -> None:
dti = pd.date_range("1/1/2025", "2/28/2025")
Expand Down