Skip to content

Commit 746c760

Browse files
cmp0xffcmp0xff
committed
setitem with none
1 parent d54943e commit 746c760

File tree

4 files changed

+56
-164
lines changed

4 files changed

+56
-164
lines changed

‎pandas-stubs/_typing.pyi‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ from pandas._libs.tslibs import (
5252
Timedelta,
5353
Timestamp,
5454
)
55+
from pandas._libs.tslibs.nattype import NaTType
5556

5657
from pandas.core.dtypes.dtypes import (
5758
CategoricalDtype,
@@ -134,6 +135,7 @@ _IndexIterScalar: TypeAlias = (
134135
Scalar: TypeAlias = (
135136
_IndexIterScalar | complex | np.integer | np.floating | np.complexfloating
136137
)
138+
ScalarOrNA: TypeAlias = Scalar | NAType | NaTType | None
137139
IntStrT = TypeVar("IntStrT", int, str)
138140

139141
# timestamp and timedelta convertible types

‎pandas-stubs/core/frame.pyi‎

Lines changed: 39 additions & 153 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,6 @@ import xarray as xr
8181
from pandas._libs.lib import _NoDefaultDoNotUse
8282
from pandas._libs.missing import NAType
8383
from pandas._libs.tslibs import BaseOffset
84-
from pandas._libs.tslibs.nattype import NaTType
8584
from pandas._typing import (
8685
S2,
8786
AggFuncTypeBase,
@@ -146,6 +145,7 @@ from pandas._typing import (
146145
Renamer,
147146
ReplaceValue,
148147
Scalar,
148+
ScalarOrNA,
149149
ScalarT,
150150
SequenceNotStr,
151151
SeriesByT,
@@ -181,6 +181,26 @@ _T_MUTABLE_MAPPING_co = TypeVar(
181181
"_T_MUTABLE_MAPPING_co", bound=MutableMapping, covariant=True
182182
)
183183

184+
_iLocSetItemKey: TypeAlias = (
185+
int
186+
| IndexType
187+
| tuple[int, int]
188+
| tuple[IndexType, int]
189+
| tuple[IndexType, IndexType]
190+
| tuple[int, IndexType]
191+
)
192+
_LocSetItemKey: TypeAlias = (
193+
MaskType | Hashable | _IndexSliceTuple | Iterable[Scalar] | IndexingInt | slice
194+
)
195+
_SetItemValueNotDataFrame: TypeAlias = (
196+
ScalarOrNA
197+
| Sequence[ScalarOrNA]
198+
| Sequence[Sequence[ScalarOrNA]]
199+
| Mapping[Any, ScalarOrNA]
200+
| ArrayLike
201+
| IndexOpsMixin
202+
)
203+
184204
class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
185205
@overload
186206
def __getitem__(self, key: tuple[int, int]) -> Scalar: ...
@@ -203,26 +223,7 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
203223

204224
# Keep in sync with `DataFrame.__setitem__`
205225
def __setitem__(
206-
self,
207-
key: (
208-
int
209-
| IndexType
210-
| tuple[int, int]
211-
| tuple[IndexType, int]
212-
| tuple[IndexType, IndexType]
213-
| tuple[int, IndexType]
214-
),
215-
value: (
216-
Scalar
217-
| IndexOpsMixin
218-
| Sequence[Scalar]
219-
| DataFrame
220-
| np_ndarray
221-
| NAType
222-
| NaTType
223-
| Mapping[Hashable, Scalar | NAType | NaTType]
224-
| None
225-
),
226+
self, key: _iLocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
226227
) -> None: ...
227228

228229
class _LocIndexerFrame(_LocIndexer, Generic[_T]):
@@ -283,52 +284,16 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
283284
# Keep in sync with `DataFrame.__setitem__`
284285
@overload
285286
def __setitem__(
286-
self,
287-
key: tuple[_IndexSliceTuple, Hashable],
288-
value: (
289-
Scalar
290-
| NAType
291-
| NaTType
292-
| ArrayLike
293-
| IndexOpsMixin
294-
| Sequence[Scalar]
295-
| Sequence[Sequence[Scalar]]
296-
| Mapping[Hashable, Scalar | NAType | NaTType]
297-
| None
298-
),
287+
self, key: tuple[_IndexSliceTuple, Hashable], value: _SetItemValueNotDataFrame
299288
) -> None: ...
300289
@overload
301290
def __setitem__(
302-
self,
303-
key: (
304-
MaskType
305-
| Hashable
306-
| _IndexSliceTuple
307-
| Iterable[Scalar]
308-
| IndexingInt
309-
| slice
310-
),
311-
value: (
312-
Scalar
313-
| NAType
314-
| NaTType
315-
| ArrayLike
316-
| IndexOpsMixin
317-
| Sequence[Scalar]
318-
| Sequence[Sequence[Scalar]]
319-
| DataFrame
320-
| Mapping[Hashable, Scalar | NAType | NaTType]
321-
| None
322-
),
291+
self, key: _LocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
323292
) -> None: ...
324293

325294
class _iAtIndexerFrame(_iAtIndexer):
326295
def __getitem__(self, key: tuple[int, int]) -> Scalar: ...
327-
def __setitem__(
328-
self,
329-
key: tuple[int, int],
330-
value: Scalar | NAType | NaTType | None,
331-
) -> None: ...
296+
def __setitem__(self, key: tuple[int, int], value: ScalarOrNA) -> None: ...
332297

333298
class _AtIndexerFrame(_AtIndexer):
334299
def __getitem__(
@@ -347,42 +312,26 @@ class _AtIndexerFrame(_AtIndexer):
347312
key: (
348313
MaskType | StrLike | _IndexSliceTuple | list[ScalarT] | IndexingInt | slice
349314
),
350-
value: (
351-
Scalar
352-
| NAType
353-
| NaTType
354-
| ArrayLike
355-
| IndexOpsMixin
356-
| DataFrame
357-
| Sequence[Scalar]
358-
| Sequence[Sequence[Scalar]]
359-
| Mapping[Hashable, Scalar | NAType | NaTType]
360-
| None
361-
),
315+
value: _SetItemValueNotDataFrame | DataFrame,
362316
) -> None: ...
363317

364-
# With mypy 1.14.1 and python 3.12, the second overload needs a type-ignore statement
365-
if sys.version_info >= (3, 12):
366-
class _GetItemHack:
367-
@overload
368-
def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
318+
# With python 3.12+, the second overload needs a type-ignore statement
319+
class _GetItemHack:
320+
@overload
321+
def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
322+
if sys.version_info >= (3, 12):
369323
@overload
370324
def __getitem__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
371325
self, key: Iterable[Hashable] | slice
372326
) -> Self: ...
373-
@overload
374-
def __getitem__(self, key: Hashable) -> Series: ...
375-
376-
else:
377-
class _GetItemHack:
378-
@overload
379-
def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
327+
else:
380328
@overload
381329
def __getitem__( # pyright: ignore[reportOverlappingOverload]
382330
self, key: Iterable[Hashable] | slice
383331
) -> Self: ...
384-
@overload
385-
def __getitem__(self, key: Hashable) -> Series: ...
332+
333+
@overload
334+
def __getitem__(self, key: Hashable) -> Series: ...
386335

387336
_AstypeArgExt: TypeAlias = (
388337
AstypeArg
@@ -829,85 +778,22 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
829778
# Keep in sync with `_iLocIndexerFrame.__setitem__`
830779
@overload
831780
def __setitem__(
832-
self,
833-
idx: (
834-
int
835-
| IndexType
836-
| tuple[int, int]
837-
| tuple[IndexType, int]
838-
| tuple[IndexType, IndexType]
839-
| tuple[int, IndexType]
840-
),
841-
value: (
842-
Scalar
843-
| IndexOpsMixin
844-
| Sequence[Scalar]
845-
| DataFrame
846-
| np_ndarray
847-
| NAType
848-
| NaTType
849-
| Mapping[Hashable, Scalar | NAType | NaTType]
850-
| None
851-
),
781+
self, idx: _iLocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
852782
) -> None: ...
853783
# Keep in sync with `_LocIndexerFrame.__setitem__`
854784
@overload
855785
def __setitem__(
856-
self,
857-
idx: tuple[_IndexSliceTuple, Hashable],
858-
value: (
859-
Scalar
860-
| NAType
861-
| NaTType
862-
| ArrayLike
863-
| IndexOpsMixin
864-
| Sequence[Scalar]
865-
| Sequence[Sequence[Scalar]]
866-
| Mapping[Hashable, Scalar | NAType | NaTType]
867-
| None
868-
),
786+
self, idx: tuple[_IndexSliceTuple, Hashable], value: _SetItemValueNotDataFrame
869787
) -> None: ...
870788
@overload
871789
def __setitem__(
872-
self,
873-
idx: (
874-
MaskType
875-
| Hashable
876-
| _IndexSliceTuple
877-
| Iterable[Scalar]
878-
| IndexingInt
879-
| slice
880-
),
881-
value: (
882-
Scalar
883-
| NAType
884-
| NaTType
885-
| ArrayLike
886-
| IndexOpsMixin
887-
| Sequence[Scalar]
888-
| Sequence[Sequence[Scalar]]
889-
| DataFrame
890-
| Mapping[Hashable, Scalar | NAType | NaTType]
891-
| None
892-
),
790+
self, idx: _LocSetItemKey, value: _SetItemValueNotDataFrame | DataFrame
893791
) -> None: ...
894792
# Extra cases not supported by `_LocIndexerFrame.__setitem__` /
895793
# `_iLocIndexerFrame.__setitem__`.
896794
@overload
897795
def __setitem__(
898-
self,
899-
idx: IndexOpsMixin | DataFrame,
900-
value: (
901-
Scalar
902-
| NAType
903-
| NaTType
904-
| ArrayLike
905-
| IndexOpsMixin
906-
| Sequence[Scalar]
907-
| Sequence[Sequence[Scalar]]
908-
| Mapping[Hashable, Scalar | NAType | NaTType]
909-
| None
910-
),
796+
self, idx: IndexOpsMixin | DataFrame, value: _SetItemValueNotDataFrame
911797
) -> None: ...
912798
@overload
913799
def query(

‎tests/frame/test_indexing.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -365,17 +365,6 @@ def test_isetframe() -> None:
365365
check(assert_type(frame.isetitem([0], [10, 12]), None), type(None))
366366

367367

368-
def test_setitem_none() -> None:
369-
df = pd.DataFrame(
370-
{"A": [1, 2, 3], "B": ["abc", "def", "ghi"]}, index=["x", "y", "z"]
371-
)
372-
df.loc["x", "B"] = None
373-
df.iloc[2, 0] = None
374-
sb = pd.Series([1, 2, 3], dtype=int)
375-
sb.loc["y"] = None
376-
sb.iloc[0] = None
377-
378-
379368
def test_getsetitem_multiindex() -> None:
380369
# GH 466
381370
rows = pd.Index(["project A", "project B", "project C"])
@@ -422,6 +411,14 @@ def test_frame_setitem_na() -> None:
422411
df.loc[ind, :] = pd.NaT
423412
df.iloc[[0, 2], :] = pd.NaT
424413

414+
df.loc["a", "x"] = None
415+
df.iloc[2, 0] = None
416+
417+
df.loc[:, "x"] = [None, pd.NA, pd.NaT]
418+
df.iloc[:, 0] = [None, pd.NA, pd.NaT]
419+
df.loc[:, ["x"]] = [[None], [pd.NA], [pd.NaT]] # type: ignore[assignment,index]
420+
df.iloc[:, [0]] = [[None], [pd.NA], [pd.NaT]] # type: ignore[assignment]
421+
425422

426423
def test_loc_set() -> None:
427424
df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
@@ -571,6 +568,9 @@ def test_df_loc_dict() -> None:
571568
df.iloc[0] = {"X": 0}
572569
check(assert_type(df, pd.DataFrame), pd.DataFrame)
573570

571+
df.loc[0] = {None: None, pd.NA: pd.NA, pd.NaT: pd.NaT}
572+
df.iloc[0] = {None: None, pd.NA: pd.NA, pd.NaT: pd.NaT}
573+
574574

575575
def test_iloc_npint() -> None:
576576
# GH 69

‎tests/series/test_indexing.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,10 @@ def test_series_setitem_na() -> None:
220220
s2.loc[ind] = pd.NaT
221221
s2.iloc[[0, 2]] = pd.NaT
222222

223+
sb = pd.Series([1, 2, 3], dtype=int)
224+
sb.loc["y"] = None
225+
sb.iloc[0] = None
226+
223227

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

0 commit comments

Comments
 (0)