Skip to content

Commit fc331b1

Browse files
loicdiridollouloicdiridollou
authored
GH984 Add overload for DataFrame.clip and update those for Series.clip (#1206)
* GH984 Add overload for DataFrame.clip and update those for Series.clip * GH984 PR feedback * GH984 drop duplicate tests * GH984 PR feedback
1 parent f5fd35f commit fc331b1

File tree

4 files changed

+401
-25
lines changed

4 files changed

+401
-25
lines changed

‎pandas-stubs/core/frame.pyi

+57-7
Original file line numberDiff line numberDiff line change
@@ -1757,23 +1757,73 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
17571757
@overload
17581758
def clip(
17591759
self,
1760-
lower: float | AnyArrayLike | None = ...,
1761-
upper: float | AnyArrayLike | None = ...,
1760+
lower: float | None = ...,
1761+
upper: float | None = ...,
17621762
*,
17631763
axis: Axis | None = ...,
1764-
inplace: Literal[True],
1764+
inplace: Literal[False] = ...,
17651765
**kwargs: Any,
1766-
) -> None: ...
1766+
) -> Self: ...
17671767
@overload
17681768
def clip(
17691769
self,
1770-
lower: float | AnyArrayLike | None = ...,
1771-
upper: float | AnyArrayLike | None = ...,
1770+
lower: AnyArrayLike = ...,
1771+
upper: AnyArrayLike | None = ...,
17721772
*,
1773-
axis: Axis | None = ...,
1773+
axis: Axis = ...,
17741774
inplace: Literal[False] = ...,
17751775
**kwargs: Any,
17761776
) -> Self: ...
1777+
@overload
1778+
def clip(
1779+
self,
1780+
lower: AnyArrayLike | None = ...,
1781+
upper: AnyArrayLike = ...,
1782+
*,
1783+
axis: Axis = ...,
1784+
inplace: Literal[False] = ...,
1785+
**kwargs: Any,
1786+
) -> Self: ...
1787+
@overload
1788+
def clip( # pyright: ignore[reportOverlappingOverload]
1789+
self,
1790+
lower: None = ...,
1791+
upper: None = ...,
1792+
*,
1793+
axis: Axis | None = ...,
1794+
inplace: Literal[True],
1795+
**kwargs: Any,
1796+
) -> Self: ...
1797+
@overload
1798+
def clip(
1799+
self,
1800+
lower: float | None = ...,
1801+
upper: float | None = ...,
1802+
*,
1803+
axis: Axis | None = ...,
1804+
inplace: Literal[True],
1805+
**kwargs: Any,
1806+
) -> None: ...
1807+
@overload
1808+
def clip(
1809+
self,
1810+
lower: AnyArrayLike = ...,
1811+
upper: AnyArrayLike | None = ...,
1812+
*,
1813+
axis: Axis = ...,
1814+
inplace: Literal[True],
1815+
**kwargs: Any,
1816+
) -> None: ...
1817+
@overload
1818+
def clip(
1819+
self,
1820+
lower: AnyArrayLike | None = ...,
1821+
upper: AnyArrayLike = ...,
1822+
*,
1823+
axis: Axis = ...,
1824+
inplace: Literal[True],
1825+
**kwargs: Any,
1826+
) -> None: ...
17771827
def copy(self, deep: _bool = ...) -> Self: ...
17781828
def cummax(
17791829
self, axis: Axis | None = ..., skipna: _bool = ..., *args: Any, **kwargs: Any

‎pandas-stubs/core/series.pyi

+10
Original file line numberDiff line numberDiff line change
@@ -1394,6 +1394,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
13941394
subset: _str | Sequence[_str] | None = ...,
13951395
) -> Scalar | Series[S1]: ...
13961396
@overload
1397+
def clip( # pyright: ignore[reportOverlappingOverload]
1398+
self,
1399+
lower: None = ...,
1400+
upper: None = ...,
1401+
*,
1402+
axis: AxisIndex | None = ...,
1403+
inplace: Literal[True],
1404+
**kwargs: Any,
1405+
) -> Self: ...
1406+
@overload
13971407
def clip(
13981408
self,
13991409
lower: AnyArrayLike | float | None = ...,

‎tests/test_frame.py

+211-16
Original file line numberDiff line numberDiff line change
@@ -659,23 +659,218 @@ def test_types_quantile() -> None:
659659
df.quantile(np.array([0.25, 0.75]))
660660

661661

662-
@pytest.mark.parametrize("lower", [None, 5, pd.Series([3, 4])])
663-
@pytest.mark.parametrize("upper", [None, 15, pd.Series([12, 13])])
664-
@pytest.mark.parametrize("axis", [None, 0, "index"])
665-
def test_types_clip(lower, upper, axis) -> None:
666-
def is_none_or_numeric(val: Any) -> bool:
667-
return val is None or isinstance(val, int | float)
668-
662+
def test_dataframe_clip() -> None:
663+
"""Test different clipping combinations for dataframe."""
669664
df = pd.DataFrame(data={"col1": [20, 12], "col2": [3, 14]})
670-
uses_array = not (is_none_or_numeric(lower) and is_none_or_numeric(upper))
671-
if uses_array and axis is None:
672-
with pytest.raises(ValueError):
673-
df.clip(lower=lower, upper=upper, axis=axis)
674-
else:
675-
check(
676-
assert_type(df.clip(lower=lower, upper=upper, axis=axis), pd.DataFrame),
677-
pd.DataFrame,
678-
)
665+
if TYPE_CHECKING_INVALID_USAGE:
666+
df.clip(lower=pd.Series([4, 5]), upper=None, axis=None) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
667+
df.clip(lower=None, upper=pd.Series([4, 5]), axis=None) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
668+
df.clip(lower=pd.Series([1, 2]), upper=pd.Series([4, 5]), axis=None) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
669+
df.copy().clip(lower=pd.Series([1, 2]), upper=None, axis=None, inplace=True) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
670+
df.copy().clip(lower=None, upper=pd.Series([1, 2]), axis=None, inplace=True) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
671+
df.copy().clip(lower=pd.Series([4, 5]), upper=pd.Series([1, 2]), axis=None, inplace=True) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
672+
673+
check(
674+
assert_type(df.clip(lower=None, upper=None, axis=None), pd.DataFrame),
675+
pd.DataFrame,
676+
)
677+
check(
678+
assert_type(df.clip(lower=5, upper=None, axis=None), pd.DataFrame), pd.DataFrame
679+
)
680+
check(
681+
assert_type(df.clip(lower=None, upper=15, axis=None), pd.DataFrame),
682+
pd.DataFrame,
683+
)
684+
check(
685+
assert_type(
686+
df.clip(lower=None, upper=None, axis=None, inplace=True), pd.DataFrame
687+
),
688+
pd.DataFrame,
689+
)
690+
check(
691+
assert_type(df.clip(lower=5, upper=None, axis=None, inplace=True), None),
692+
type(None),
693+
)
694+
check(
695+
assert_type(df.clip(lower=None, upper=15, axis=None, inplace=True), None),
696+
type(None),
697+
)
698+
699+
check(
700+
assert_type(df.clip(lower=None, upper=None, axis=0), pd.DataFrame), pd.DataFrame
701+
)
702+
check(assert_type(df.clip(lower=5, upper=None, axis=0), pd.DataFrame), pd.DataFrame)
703+
check(
704+
assert_type(df.clip(lower=None, upper=15, axis=0), pd.DataFrame), pd.DataFrame
705+
)
706+
check(
707+
assert_type(df.clip(lower=pd.Series([1, 2]), upper=None, axis=0), pd.DataFrame),
708+
pd.DataFrame,
709+
)
710+
check(
711+
assert_type(df.clip(lower=None, upper=pd.Series([1, 2]), axis=0), pd.DataFrame),
712+
pd.DataFrame,
713+
)
714+
check(
715+
assert_type(
716+
df.clip(lower=None, upper=None, axis="index", inplace=True), pd.DataFrame
717+
),
718+
pd.DataFrame,
719+
)
720+
check(
721+
assert_type(df.clip(lower=5, upper=None, axis="index", inplace=True), None),
722+
type(None),
723+
)
724+
check(
725+
assert_type(df.clip(lower=None, upper=15, axis="index", inplace=True), None),
726+
type(None),
727+
)
728+
check(
729+
assert_type(
730+
df.clip(lower=pd.Series([1, 2]), upper=None, axis="index", inplace=True),
731+
None,
732+
),
733+
type(None),
734+
)
735+
check(
736+
assert_type(
737+
df.clip(lower=None, upper=pd.Series([1, 2]), axis="index", inplace=True),
738+
None,
739+
),
740+
type(None),
741+
)
742+
check(
743+
assert_type(df.clip(lower=None, upper=None, axis="index"), pd.DataFrame),
744+
pd.DataFrame,
745+
)
746+
check(
747+
assert_type(df.clip(lower=5, upper=None, axis="index"), pd.DataFrame),
748+
pd.DataFrame,
749+
)
750+
check(
751+
assert_type(df.clip(lower=None, upper=15, axis="index"), pd.DataFrame),
752+
pd.DataFrame,
753+
)
754+
check(
755+
assert_type(
756+
df.clip(lower=pd.Series([1, 2]), upper=None, axis="index"), pd.DataFrame
757+
),
758+
pd.DataFrame,
759+
)
760+
check(
761+
assert_type(
762+
df.clip(lower=None, upper=pd.Series([1, 2]), axis="index"), pd.DataFrame
763+
),
764+
pd.DataFrame,
765+
)
766+
check(
767+
assert_type(
768+
df.clip(lower=None, upper=None, axis=0, inplace=True), pd.DataFrame
769+
),
770+
pd.DataFrame,
771+
)
772+
check(
773+
assert_type(df.clip(lower=5, upper=None, axis=0, inplace=True), None),
774+
type(None),
775+
)
776+
check(
777+
assert_type(df.clip(lower=None, upper=15, axis=0, inplace=True), None),
778+
type(None),
779+
)
780+
781+
# without lower
782+
check(assert_type(df.clip(upper=None, axis=None), pd.DataFrame), pd.DataFrame)
783+
check(assert_type(df.clip(upper=15, axis=None), pd.DataFrame), pd.DataFrame)
784+
check(
785+
assert_type(df.clip(upper=None, axis=None, inplace=True), pd.DataFrame),
786+
pd.DataFrame,
787+
)
788+
check(assert_type(df.clip(upper=15, axis=None, inplace=True), None), type(None))
789+
790+
check(assert_type(df.clip(upper=None, axis=0), pd.DataFrame), pd.DataFrame)
791+
check(assert_type(df.clip(upper=15, axis=0), pd.DataFrame), pd.DataFrame)
792+
check(
793+
assert_type(df.clip(upper=pd.Series([1, 2]), axis=0), pd.DataFrame),
794+
pd.DataFrame,
795+
)
796+
check(
797+
assert_type(df.clip(upper=None, axis="index", inplace=True), pd.DataFrame),
798+
pd.DataFrame,
799+
)
800+
check(
801+
assert_type(df.clip(upper=15, axis="index", inplace=True), None),
802+
type(None),
803+
)
804+
check(
805+
assert_type(df.clip(upper=pd.Series([1, 2]), axis="index", inplace=True), None),
806+
type(None),
807+
)
808+
check(assert_type(df.clip(upper=None, axis="index"), pd.DataFrame), pd.DataFrame)
809+
check(assert_type(df.clip(upper=15, axis="index"), pd.DataFrame), pd.DataFrame)
810+
check(
811+
assert_type(df.clip(upper=pd.Series([1, 2]), axis="index"), pd.DataFrame),
812+
pd.DataFrame,
813+
)
814+
check(
815+
assert_type(df.clip(upper=None, axis=0, inplace=True), pd.DataFrame),
816+
pd.DataFrame,
817+
)
818+
check(assert_type(df.clip(upper=15, axis=0, inplace=True), None), type(None))
819+
820+
# without upper
821+
check(
822+
assert_type(df.clip(lower=None, axis=None), pd.DataFrame),
823+
pd.DataFrame,
824+
)
825+
check(assert_type(df.clip(lower=5, axis=None), pd.DataFrame), pd.DataFrame)
826+
check(
827+
assert_type(df.clip(lower=None, axis=None, inplace=True), pd.DataFrame),
828+
pd.DataFrame,
829+
)
830+
check(
831+
assert_type(df.clip(lower=5, axis=None, inplace=True), None),
832+
type(None),
833+
)
834+
check(
835+
assert_type(df.clip(lower=None, axis=None, inplace=True), pd.DataFrame),
836+
pd.DataFrame,
837+
)
838+
839+
check(assert_type(df.clip(lower=None, axis=0), pd.DataFrame), pd.DataFrame)
840+
check(assert_type(df.clip(lower=5, axis=0), pd.DataFrame), pd.DataFrame)
841+
check(
842+
assert_type(df.clip(lower=pd.Series([1, 2]), axis=0), pd.DataFrame),
843+
pd.DataFrame,
844+
)
845+
check(
846+
assert_type(df.clip(lower=None, axis="index", inplace=True), pd.DataFrame),
847+
pd.DataFrame,
848+
)
849+
check(
850+
assert_type(df.clip(lower=5, axis="index", inplace=True), None),
851+
type(None),
852+
)
853+
check(
854+
assert_type(df.clip(lower=pd.Series([1, 2]), axis="index", inplace=True), None),
855+
type(None),
856+
)
857+
check(
858+
assert_type(df.clip(lower=None, axis="index"), pd.DataFrame),
859+
pd.DataFrame,
860+
)
861+
check(
862+
assert_type(df.clip(lower=5, axis="index"), pd.DataFrame),
863+
pd.DataFrame,
864+
)
865+
check(
866+
assert_type(df.clip(lower=pd.Series([1, 2]), axis="index"), pd.DataFrame),
867+
pd.DataFrame,
868+
)
869+
check(
870+
assert_type(df.clip(lower=None, axis=0, inplace=True), pd.DataFrame),
871+
pd.DataFrame,
872+
)
873+
check(assert_type(df.clip(lower=5, axis=0, inplace=True), None), type(None))
679874

680875

681876
def test_types_abs() -> None:

0 commit comments

Comments
 (0)