Skip to content

Commit 8621c71

Browse files
JanEricNitschkeJanEricNitschke
committed
Encode more argument restrictions in the signature overloads.
1 parent 97b2655 commit 8621c71

File tree

5 files changed

+253
-9
lines changed

5 files changed

+253
-9
lines changed

‎pandas-stubs/core/frame.pyi

Lines changed: 75 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -697,10 +697,22 @@ class DataFrame(NDFrame, OpsMixin):
697697
@overload
698698
def drop(
699699
self,
700-
labels: Hashable | Sequence[Hashable] | Index[Any] = ...,
700+
labels: None = ...,
701701
*,
702702
axis: Axis = ...,
703703
index: Hashable | Sequence[Hashable] | Index[Any] = ...,
704+
columns: Hashable | Sequence[Hashable] | Index[Any],
705+
level: Level | None = ...,
706+
inplace: Literal[True],
707+
errors: IgnoreRaise = ...,
708+
) -> None: ...
709+
@overload
710+
def drop(
711+
self,
712+
labels: None = ...,
713+
*,
714+
axis: Axis = ...,
715+
index: Hashable | Sequence[Hashable] | Index[Any],
704716
columns: Hashable | Sequence[Hashable] | Index[Any] = ...,
705717
level: Level | None = ...,
706718
inplace: Literal[True],
@@ -709,10 +721,34 @@ class DataFrame(NDFrame, OpsMixin):
709721
@overload
710722
def drop(
711723
self,
712-
labels: Hashable | Sequence[Hashable] | Index[Any] = ...,
724+
labels: Hashable | Sequence[Hashable] | Index[Any],
725+
*,
726+
axis: Axis = ...,
727+
index: None = ...,
728+
columns: None = ...,
729+
level: Level | None = ...,
730+
inplace: Literal[True],
731+
errors: IgnoreRaise = ...,
732+
) -> None: ...
733+
@overload
734+
def drop(
735+
self,
736+
labels: None = ...,
713737
*,
714738
axis: Axis = ...,
715739
index: Hashable | Sequence[Hashable] | Index[Any] = ...,
740+
columns: Hashable | Sequence[Hashable] | Index[Any],
741+
level: Level | None = ...,
742+
inplace: Literal[False] = ...,
743+
errors: IgnoreRaise = ...,
744+
) -> DataFrame: ...
745+
@overload
746+
def drop(
747+
self,
748+
labels: None = ...,
749+
*,
750+
axis: Axis = ...,
751+
index: Hashable | Sequence[Hashable] | Index[Any],
716752
columns: Hashable | Sequence[Hashable] | Index[Any] = ...,
717753
level: Level | None = ...,
718754
inplace: Literal[False] = ...,
@@ -721,16 +757,52 @@ class DataFrame(NDFrame, OpsMixin):
721757
@overload
722758
def drop(
723759
self,
724-
labels: Hashable | Sequence[Hashable] | Index[Any] = ...,
760+
labels: Hashable | Sequence[Hashable] | Index[Any],
761+
*,
762+
axis: Axis = ...,
763+
index: None = ...,
764+
columns: None = ...,
765+
level: Level | None = ...,
766+
inplace: Literal[False] = ...,
767+
errors: IgnoreRaise = ...,
768+
) -> DataFrame: ...
769+
@overload
770+
def drop(
771+
self,
772+
labels: None = ...,
725773
*,
726774
axis: Axis = ...,
727775
index: Hashable | Sequence[Hashable] | Index[Any] = ...,
776+
columns: Hashable | Sequence[Hashable] | Index[Any],
777+
level: Level | None = ...,
778+
inplace: bool = ...,
779+
errors: IgnoreRaise = ...,
780+
) -> DataFrame | None: ...
781+
@overload
782+
def drop(
783+
self,
784+
labels: None = ...,
785+
*,
786+
axis: Axis = ...,
787+
index: Hashable | Sequence[Hashable] | Index[Any],
728788
columns: Hashable | Sequence[Hashable] | Index[Any] = ...,
729789
level: Level | None = ...,
730790
inplace: bool = ...,
731791
errors: IgnoreRaise = ...,
732792
) -> DataFrame | None: ...
733793
@overload
794+
def drop(
795+
self,
796+
labels: Hashable | Sequence[Hashable] | Index[Any],
797+
*,
798+
axis: Axis = ...,
799+
index: None = ...,
800+
columns: None = ...,
801+
level: Level | None = ...,
802+
inplace: bool = ...,
803+
errors: IgnoreRaise = ...,
804+
) -> DataFrame | None: ...
805+
@overload
734806
def rename(
735807
self,
736808
mapper: Renamer | None = ...,

‎pandas-stubs/core/generic.pyi

Lines changed: 75 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -305,10 +305,22 @@ class NDFrame(indexing.IndexingMixin):
305305
@overload
306306
def drop(
307307
self,
308-
labels: Hashable | Sequence[Hashable] = ...,
308+
labels: None = ...,
309309
*,
310310
axis: Axis = ...,
311311
index: Hashable | Sequence[Hashable] = ...,
312+
columns: Hashable | Sequence[Hashable],
313+
level: Level | None = ...,
314+
inplace: Literal[True],
315+
errors: IgnoreRaise = ...,
316+
) -> None: ...
317+
@overload
318+
def drop(
319+
self,
320+
labels: None = ...,
321+
*,
322+
axis: Axis = ...,
323+
index: Hashable | Sequence[Hashable],
312324
columns: Hashable | Sequence[Hashable] = ...,
313325
level: Level | None = ...,
314326
inplace: Literal[True],
@@ -317,10 +329,34 @@ class NDFrame(indexing.IndexingMixin):
317329
@overload
318330
def drop(
319331
self,
320-
labels: Hashable | Sequence[Hashable] = ...,
332+
labels: Hashable | Sequence[Hashable],
333+
*,
334+
axis: Axis = ...,
335+
index: None = ...,
336+
columns: None = ...,
337+
level: Level | None = ...,
338+
inplace: Literal[True],
339+
errors: IgnoreRaise = ...,
340+
) -> None: ...
341+
@overload
342+
def drop(
343+
self,
344+
labels: None = ...,
321345
*,
322346
axis: Axis = ...,
323347
index: Hashable | Sequence[Hashable] = ...,
348+
columns: Hashable | Sequence[Hashable],
349+
level: Level | None = ...,
350+
inplace: Literal[False] = ...,
351+
errors: IgnoreRaise = ...,
352+
) -> Self: ...
353+
@overload
354+
def drop(
355+
self,
356+
labels: None = ...,
357+
*,
358+
axis: Axis = ...,
359+
index: Hashable | Sequence[Hashable],
324360
columns: Hashable | Sequence[Hashable] = ...,
325361
level: Level | None = ...,
326362
inplace: Literal[False] = ...,
@@ -329,15 +365,51 @@ class NDFrame(indexing.IndexingMixin):
329365
@overload
330366
def drop(
331367
self,
332-
labels: Hashable | Sequence[Hashable] = ...,
368+
labels: Hashable | Sequence[Hashable],
369+
*,
370+
axis: Axis = ...,
371+
index: None = ...,
372+
columns: None = ...,
373+
level: Level | None = ...,
374+
inplace: Literal[False] = ...,
375+
errors: IgnoreRaise = ...,
376+
) -> Self: ...
377+
@overload
378+
def drop(
379+
self,
380+
labels: None = ...,
333381
*,
334382
axis: Axis = ...,
335383
index: Hashable | Sequence[Hashable] = ...,
384+
columns: Hashable | Sequence[Hashable],
385+
level: Level | None = ...,
386+
inplace: _bool = ...,
387+
errors: IgnoreRaise = ...,
388+
) -> Self | None: ...
389+
@overload
390+
def drop(
391+
self,
392+
labels: None = ...,
393+
*,
394+
axis: Axis = ...,
395+
index: Hashable | Sequence[Hashable],
336396
columns: Hashable | Sequence[Hashable] = ...,
337397
level: Level | None = ...,
338398
inplace: _bool = ...,
339399
errors: IgnoreRaise = ...,
340400
) -> Self | None: ...
401+
@overload
402+
def drop(
403+
self,
404+
labels: Hashable | Sequence[Hashable],
405+
*,
406+
axis: Axis = ...,
407+
index: None = ...,
408+
columns: None = ...,
409+
level: Level | None = ...,
410+
inplace: _bool = ...,
411+
errors: IgnoreRaise = ...,
412+
) -> Self | None: ...
341413
def add_prefix(self, prefix: _str) -> Self: ...
342414
def add_suffix(self, suffix: _str) -> Self: ...
343415
def sort_index(

‎pandas-stubs/core/series.pyi

Lines changed: 75 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1054,10 +1054,22 @@ class Series(IndexOpsMixin[S1], NDFrame):
10541054
@overload
10551055
def drop(
10561056
self,
1057-
labels: Hashable | list[HashableT1] | Index = ...,
1057+
labels: None = ...,
10581058
*,
10591059
axis: Axis = ...,
10601060
index: Hashable | list[HashableT2] | Index = ...,
1061+
columns: Hashable | list[HashableT3] | Index,
1062+
level: Level | None = ...,
1063+
inplace: Literal[True],
1064+
errors: IgnoreRaise = ...,
1065+
) -> None: ...
1066+
@overload
1067+
def drop(
1068+
self,
1069+
labels: None = ...,
1070+
*,
1071+
axis: Axis = ...,
1072+
index: Hashable | list[HashableT2] | Index,
10611073
columns: Hashable | list[HashableT3] | Index = ...,
10621074
level: Level | None = ...,
10631075
inplace: Literal[True],
@@ -1066,10 +1078,34 @@ class Series(IndexOpsMixin[S1], NDFrame):
10661078
@overload
10671079
def drop(
10681080
self,
1069-
labels: Hashable | list[HashableT1] | Index = ...,
1081+
labels: Hashable | list[HashableT1] | Index,
1082+
*,
1083+
axis: Axis = ...,
1084+
index: None = ...,
1085+
columns: None = ...,
1086+
level: Level | None = ...,
1087+
inplace: Literal[True],
1088+
errors: IgnoreRaise = ...,
1089+
) -> None: ...
1090+
@overload
1091+
def drop(
1092+
self,
1093+
labels: None = ...,
10701094
*,
10711095
axis: Axis = ...,
10721096
index: Hashable | list[HashableT2] | Index = ...,
1097+
columns: Hashable | list[HashableT3] | Index,
1098+
level: Level | None = ...,
1099+
inplace: Literal[False] = ...,
1100+
errors: IgnoreRaise = ...,
1101+
) -> Self: ...
1102+
@overload
1103+
def drop(
1104+
self,
1105+
labels: None = ...,
1106+
*,
1107+
axis: Axis = ...,
1108+
index: Hashable | list[HashableT2] | Index,
10731109
columns: Hashable | list[HashableT3] | Index = ...,
10741110
level: Level | None = ...,
10751111
inplace: Literal[False] = ...,
@@ -1078,16 +1114,52 @@ class Series(IndexOpsMixin[S1], NDFrame):
10781114
@overload
10791115
def drop(
10801116
self,
1081-
labels: Hashable | list[HashableT1] | Index = ...,
1117+
labels: Hashable | list[HashableT1] | Index,
1118+
*,
1119+
axis: Axis = ...,
1120+
index: None = ...,
1121+
columns: None = ...,
1122+
level: Level | None = ...,
1123+
inplace: Literal[False] = ...,
1124+
errors: IgnoreRaise = ...,
1125+
) -> Self: ...
1126+
@overload
1127+
def drop(
1128+
self,
1129+
labels: None = ...,
10821130
*,
10831131
axis: Axis = ...,
10841132
index: Hashable | list[HashableT2] | Index = ...,
1133+
columns: Hashable | list[HashableT3] | Index,
1134+
level: Level | None = ...,
1135+
inplace: bool = ...,
1136+
errors: IgnoreRaise = ...,
1137+
) -> Series | None: ...
1138+
@overload
1139+
def drop(
1140+
self,
1141+
labels: None = ...,
1142+
*,
1143+
axis: Axis = ...,
1144+
index: Hashable | list[HashableT2] | Index,
10851145
columns: Hashable | list[HashableT3] | Index = ...,
10861146
level: Level | None = ...,
10871147
inplace: bool = ...,
10881148
errors: IgnoreRaise = ...,
10891149
) -> Series | None: ...
10901150
@overload
1151+
def drop(
1152+
self,
1153+
labels: Hashable | list[HashableT1] | Index,
1154+
*,
1155+
axis: Axis = ...,
1156+
index: None = ...,
1157+
columns: None = ...,
1158+
level: Level | None = ...,
1159+
inplace: bool = ...,
1160+
errors: IgnoreRaise = ...,
1161+
) -> Series | None: ...
1162+
@overload
10911163
def fillna(
10921164
self,
10931165
value: Scalar | NAType | dict | Series[S1] | DataFrame | None = ...,

‎tests/test_frame.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,20 @@ def test_types_drop() -> None:
344344
check(assert_type(df.drop(columns=pd.Index(["col1"])), pd.DataFrame), pd.DataFrame)
345345

346346

347+
def test_arguments_drop() -> None:
348+
# GH 950
349+
if TYPE_CHECKING_INVALID_USAGE:
350+
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
351+
res1 = df.drop() # type: ignore[call-overload] # pyright: ignore[reportCallIssue]
352+
res2 = df.drop([0], columns=["col1"]) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
353+
res3 = df.drop([0], index=[0]) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
354+
# These should also fail, but `None` is Hasheable and i do not know how
355+
# to type hint a non-None hashable.
356+
# res4 = df.drop(columns=None)
357+
# res5 = df.drop(index=None)
358+
# res6 = df.drop(None)
359+
360+
347361
def test_types_dropna() -> None:
348362
df = pd.DataFrame(data={"col1": [np.nan, np.nan], "col2": [3, np.nan]})
349363
res: pd.DataFrame = df.dropna()

‎tests/test_series.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,20 @@ def test_types_drop() -> None:
266266
)
267267

268268

269+
def test_arguments_drop() -> None:
270+
# GH 950
271+
if TYPE_CHECKING_INVALID_USAGE:
272+
s = pd.Series([0, 1, 2])
273+
res1 = s.drop() # type: ignore[call-overload] # pyright: ignore[reportCallIssue]
274+
res2 = s.drop([0], columns=["col1"]) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
275+
res3 = s.drop([0], index=[0]) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
276+
# These should also fail, but `None` is Hasheable and i do not know how
277+
# to type hint a non-None hashable.
278+
# res4 = s.drop(columns=None)
279+
# res5 = s.drop(index=None)
280+
# res6 = s.drop(None)
281+
282+
269283
def test_types_drop_multilevel() -> None:
270284
index = pd.MultiIndex(
271285
levels=[["top", "bottom"], ["first", "second", "third"]],

0 commit comments

Comments
 (0)