Skip to content

Commit 7e5b771

Browse files
JanEricNitschkeJanEricNitschke
authored
Expand drop() overloads to encode more argument restriction and add type parameters to df.drop() signature. (#951)
* Add type parameter to df.drop() * Encode more argument restrictions in the signature overloads. * Move drop overloads from series and df into generic.
1 parent 63c4567 commit 7e5b771

File tree

5 files changed

+109
-83
lines changed

5 files changed

+109
-83
lines changed

‎pandas-stubs/core/frame.pyi‎

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -695,42 +695,6 @@ class DataFrame(NDFrame, OpsMixin):
695695
tolerance: float | None = ...,
696696
) -> DataFrame: ...
697697
@overload
698-
def drop(
699-
self,
700-
labels: Hashable | Sequence[Hashable] | Index = ...,
701-
*,
702-
axis: Axis = ...,
703-
index: Hashable | Sequence[Hashable] | Index = ...,
704-
columns: Hashable | Sequence[Hashable] | Index = ...,
705-
level: Level | None = ...,
706-
inplace: Literal[True],
707-
errors: IgnoreRaise = ...,
708-
) -> None: ...
709-
@overload
710-
def drop(
711-
self,
712-
labels: Hashable | Sequence[Hashable] | Index = ...,
713-
*,
714-
axis: Axis = ...,
715-
index: Hashable | Sequence[Hashable] | Index = ...,
716-
columns: Hashable | Sequence[Hashable] | Index = ...,
717-
level: Level | None = ...,
718-
inplace: Literal[False] = ...,
719-
errors: IgnoreRaise = ...,
720-
) -> DataFrame: ...
721-
@overload
722-
def drop(
723-
self,
724-
labels: Hashable | Sequence[Hashable] | Index = ...,
725-
*,
726-
axis: Axis = ...,
727-
index: Hashable | Sequence[Hashable] | Index = ...,
728-
columns: Hashable | Sequence[Hashable] | Index = ...,
729-
level: Level | None = ...,
730-
inplace: bool = ...,
731-
errors: IgnoreRaise = ...,
732-
) -> DataFrame | None: ...
733-
@overload
734698
def rename(
735699
self,
736700
mapper: Renamer | None = ...,

‎pandas-stubs/core/generic.pyi‎

Lines changed: 81 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -305,35 +305,107 @@ class NDFrame(indexing.IndexingMixin):
305305
@overload
306306
def drop(
307307
self,
308-
labels: Hashable | Sequence[Hashable] = ...,
308+
labels: None = ...,
309309
*,
310310
axis: Axis = ...,
311-
index: Hashable | Sequence[Hashable] = ...,
312-
columns: Hashable | Sequence[Hashable] = ...,
311+
index: Hashable | Sequence[Hashable] | Index[Any] = ...,
312+
columns: Hashable | Sequence[Hashable] | Index[Any],
313313
level: Level | None = ...,
314314
inplace: Literal[True],
315315
errors: IgnoreRaise = ...,
316316
) -> None: ...
317317
@overload
318318
def drop(
319319
self,
320-
labels: Hashable | Sequence[Hashable] = ...,
320+
labels: None = ...,
321321
*,
322322
axis: Axis = ...,
323-
index: Hashable | Sequence[Hashable] = ...,
324-
columns: Hashable | Sequence[Hashable] = ...,
323+
index: Hashable | Sequence[Hashable] | Index[Any],
324+
columns: Hashable | Sequence[Hashable] | Index[Any] = ...,
325+
level: Level | None = ...,
326+
inplace: Literal[True],
327+
errors: IgnoreRaise = ...,
328+
) -> None: ...
329+
@overload
330+
def drop(
331+
self,
332+
labels: Hashable | Sequence[Hashable] | Index[Any],
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 = ...,
345+
*,
346+
axis: Axis = ...,
347+
index: Hashable | Sequence[Hashable] | Index[Any] = ...,
348+
columns: Hashable | Sequence[Hashable] | Index[Any],
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] | Index[Any],
360+
columns: Hashable | Sequence[Hashable] | Index[Any] = ...,
325361
level: Level | None = ...,
326362
inplace: Literal[False] = ...,
327363
errors: IgnoreRaise = ...,
328364
) -> Self: ...
329365
@overload
330366
def drop(
331367
self,
332-
labels: Hashable | Sequence[Hashable] = ...,
368+
labels: Hashable | Sequence[Hashable] | Index[Any],
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 = ...,
381+
*,
382+
axis: Axis = ...,
383+
index: Hashable | Sequence[Hashable] | Index[Any] = ...,
384+
columns: Hashable | Sequence[Hashable] | Index[Any],
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] | Index[Any],
396+
columns: Hashable | Sequence[Hashable] | Index[Any] = ...,
397+
level: Level | None = ...,
398+
inplace: _bool = ...,
399+
errors: IgnoreRaise = ...,
400+
) -> Self | None: ...
401+
@overload
402+
def drop(
403+
self,
404+
labels: Hashable | Sequence[Hashable] | Index[Any],
333405
*,
334406
axis: Axis = ...,
335-
index: Hashable | Sequence[Hashable] = ...,
336-
columns: Hashable | Sequence[Hashable] = ...,
407+
index: None = ...,
408+
columns: None = ...,
337409
level: Level | None = ...,
338410
inplace: _bool = ...,
339411
errors: IgnoreRaise = ...,

‎pandas-stubs/core/series.pyi‎

Lines changed: 0 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -116,8 +116,6 @@ from pandas._typing import (
116116
FloatDtypeArg,
117117
GroupByObjectNonScalar,
118118
HashableT1,
119-
HashableT2,
120-
HashableT3,
121119
IgnoreRaise,
122120
IndexingInt,
123121
IndexLabel,
@@ -1052,42 +1050,6 @@ class Series(IndexOpsMixin[S1], NDFrame):
10521050
tolerance: float | None = ...,
10531051
) -> Self: ...
10541052
@overload
1055-
def drop(
1056-
self,
1057-
labels: Hashable | list[HashableT1] | Index = ...,
1058-
*,
1059-
axis: Axis = ...,
1060-
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: Hashable | list[HashableT1] | Index = ...,
1070-
*,
1071-
axis: Axis = ...,
1072-
index: Hashable | list[HashableT2] | Index = ...,
1073-
columns: Hashable | list[HashableT3] | Index = ...,
1074-
level: Level | None = ...,
1075-
inplace: Literal[False] = ...,
1076-
errors: IgnoreRaise = ...,
1077-
) -> Self: ...
1078-
@overload
1079-
def drop(
1080-
self,
1081-
labels: Hashable | list[HashableT1] | Index = ...,
1082-
*,
1083-
axis: Axis = ...,
1084-
index: Hashable | list[HashableT2] | Index = ...,
1085-
columns: Hashable | list[HashableT3] | Index = ...,
1086-
level: Level | None = ...,
1087-
inplace: bool = ...,
1088-
errors: IgnoreRaise = ...,
1089-
) -> Series | None: ...
1090-
@overload
10911053
def fillna(
10921054
self,
10931055
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)