Skip to content

Commit 6f392cb

Browse files
cmp0xffcmp0xff
committed
setitem with none
1 parent d54943e commit 6f392cb

File tree

2 files changed

+33
-138
lines changed

2 files changed

+33
-138
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: 31 additions & 138 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+
_iLocSetItemOneKey: TypeAlias = (
185+
int
186+
| IndexType
187+
| tuple[int, int]
188+
| tuple[IndexType, int]
189+
| tuple[IndexType, IndexType]
190+
| tuple[int, IndexType]
191+
)
192+
_LocSetItemOneKey: 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: _iLocSetItemOneKey, 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: _LocSetItemOneKey, 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,18 +312,7 @@ 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

364318
# With mypy 1.14.1 and python 3.12, the second overload needs a type-ignore statement
@@ -829,85 +783,24 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
829783
# Keep in sync with `_iLocIndexerFrame.__setitem__`
830784
@overload
831785
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-
),
786+
self, idx: _iLocSetItemOneKey, value: _SetItemValueNotDataFrame | DataFrame
852787
) -> None: ...
853788
# Keep in sync with `_LocIndexerFrame.__setitem__`
854789
@overload
855790
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-
),
791+
self, idx: tuple[_IndexSliceTuple, Hashable], value: _SetItemValueNotDataFrame
869792
) -> None: ...
870793
@overload
871794
def __setitem__(
872795
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-
),
796+
idx: _LocSetItemOneKey,
797+
value: _SetItemValueNotDataFrame | DataFrame,
893798
) -> None: ...
894799
# Extra cases not supported by `_LocIndexerFrame.__setitem__` /
895800
# `_iLocIndexerFrame.__setitem__`.
896801
@overload
897802
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-
),
803+
self, idx: IndexOpsMixin | DataFrame, value: _SetItemValueNotDataFrame
911804
) -> None: ...
912805
@overload
913806
def query(

0 commit comments

Comments
 (0)