Skip to content

Commit a8e34a4

Browse files
committed
illustrate differentiating bools
1 parent c8e3ab2 commit a8e34a4

File tree

2 files changed

+25
-7
lines changed

2 files changed

+25
-7
lines changed

pandas-stubs/core/series.pyi

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ from typing import (
2525
Generic,
2626
Literal,
2727
NoReturn,
28+
Protocol,
2829
TypeVar,
2930
final,
3031
overload,
@@ -91,6 +92,7 @@ from typing_extensions import (
9192
Never,
9293
Self,
9394
TypeAlias,
95+
override,
9496
)
9597
import xarray as xr
9698

@@ -191,6 +193,14 @@ from pandas.plotting import PlotAccessor
191193
_T_INT = TypeVar("_T_INT", bound=int)
192194
_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex)
193195

196+
class Just(Protocol, Generic[_T]):
197+
@property # type: ignore[override]
198+
@override
199+
def __class__(self, /) -> type[_T]: ...
200+
@__class__.setter
201+
@override
202+
def __class__(self, t: type[_T], /) -> None: ...
203+
194204
class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
195205
# get item
196206
@overload
@@ -2151,19 +2161,27 @@ class Series(IndexOpsMixin[S1], NDFrame):
21512161
@overload
21522162
def sub(
21532163
self: Series[bool],
2154-
other: bool | Sequence[bool],
2164+
other: Just[int] | Sequence[Just[int]],
21552165
level: Level | None = None,
21562166
fill_value: float | None = None,
21572167
axis: int = 0,
2158-
) -> Never: ...
2168+
) -> Series[int]: ...
21592169
@overload
21602170
def sub(
21612171
self: Series[bool],
2162-
other: _T_COMPLEX | Sequence[_T_COMPLEX],
2172+
other: Just[float] | Sequence[Just[float]],
21632173
level: Level | None = None,
21642174
fill_value: float | None = None,
21652175
axis: int = 0,
2166-
) -> Series[_T_COMPLEX]: ...
2176+
) -> Series[float]: ...
2177+
@overload
2178+
def sub(
2179+
self: Series[bool],
2180+
other: Just[complex] | Sequence[Just[complex]],
2181+
level: Level | None = None,
2182+
fill_value: float | None = None,
2183+
axis: int = 0,
2184+
) -> Series[complex]: ...
21672185
@overload
21682186
def sub(
21692187
self: Series[bool],

tests/series/arithmetic/bool/test_sub.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def test_sub_py_scalar() -> None:
3333
check(assert_type(c - left, "pd.Series[complex]"), pd.Series, np.complexfloating)
3434

3535
if TYPE_CHECKING_INVALID_USAGE:
36-
assert_type(left.sub(b), Never)
36+
left.sub(b) # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]
3737
check(assert_type(left.sub(i), "pd.Series[int]"), pd.Series, np.integer)
3838
check(assert_type(left.sub(f), "pd.Series[float]"), pd.Series, np.floating)
3939
check(assert_type(left.sub(c), "pd.Series[complex]"), pd.Series, np.complexfloating)
@@ -64,7 +64,7 @@ def test_sub_py_sequence() -> None:
6464
check(assert_type(c - left, "pd.Series[complex]"), pd.Series, np.complexfloating)
6565

6666
if TYPE_CHECKING_INVALID_USAGE:
67-
assert_type(left.sub(b), Never)
67+
left.sub(b) # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]
6868
check(assert_type(left.sub(i), "pd.Series[int]"), pd.Series, np.integer)
6969
check(assert_type(left.sub(f), "pd.Series[float]"), pd.Series, np.floating)
7070
check(assert_type(left.sub(c), "pd.Series[complex]"), pd.Series, np.complexfloating)
@@ -128,7 +128,7 @@ def test_sub_pd_series() -> None:
128128

129129
# In the following two cases, mypy fails to recognise the second operand as pd.Series[bool]
130130
if TYPE_CHECKING_INVALID_USAGE:
131-
_ = left - cast("pd.Series[bool]", b) # type: ignore[redundant-cast,operator] # pyright: ignore[reportUnnecessaryCast,reportOperatorIssue]
131+
result = left - b # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
132132
check(assert_type(left - i, "pd.Series[int]"), pd.Series, np.integer)
133133
check(assert_type(left - f, "pd.Series[float]"), pd.Series, np.floating)
134134
check(assert_type(left - c, "pd.Series[complex]"), pd.Series, np.complexfloating)

0 commit comments

Comments
 (0)