From cd5be5374517567b0d4d9aaf054fac1ba345bdd9 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Wed, 27 Aug 2025 22:25:26 +0200 Subject: [PATCH 1/4] feat(index): arithmetic addition --- pandas-stubs/_typing.pyi | 2 + pandas-stubs/core/indexes/base.pyi | 166 +++++++++++ pandas-stubs/core/indexes/datetimes.pyi | 4 +- pandas-stubs/core/indexes/timedeltas.pyi | 6 +- pandas-stubs/core/series.pyi | 290 +++++++++---------- tests/indexes/__init__.py | 0 tests/indexes/arithmetic/__init__.py | 0 tests/indexes/arithmetic/bool/__init__.py | 0 tests/indexes/arithmetic/bool/test_add.py | 83 ++++++ tests/indexes/arithmetic/complex/__init__.py | 0 tests/indexes/arithmetic/complex/test_add.py | 84 ++++++ tests/indexes/arithmetic/float/__init__.py | 0 tests/indexes/arithmetic/float/test_add.py | 82 ++++++ tests/indexes/arithmetic/int/__init__.py | 0 tests/indexes/arithmetic/int/test_add.py | 82 ++++++ tests/indexes/arithmetic/str/__init__.py | 0 tests/indexes/arithmetic/str/test_add.py | 88 ++++++ tests/indexes/arithmetic/test_add.py | 87 ++++++ tests/{ => indexes}/test_indexes.py | 0 19 files changed, 824 insertions(+), 150 deletions(-) create mode 100644 tests/indexes/__init__.py create mode 100644 tests/indexes/arithmetic/__init__.py create mode 100644 tests/indexes/arithmetic/bool/__init__.py create mode 100644 tests/indexes/arithmetic/bool/test_add.py create mode 100644 tests/indexes/arithmetic/complex/__init__.py create mode 100644 tests/indexes/arithmetic/complex/test_add.py create mode 100644 tests/indexes/arithmetic/float/__init__.py create mode 100644 tests/indexes/arithmetic/float/test_add.py create mode 100644 tests/indexes/arithmetic/int/__init__.py create mode 100644 tests/indexes/arithmetic/int/test_add.py create mode 100644 tests/indexes/arithmetic/str/__init__.py create mode 100644 tests/indexes/arithmetic/str/test_add.py create mode 100644 tests/indexes/arithmetic/test_add.py rename tests/{ => indexes}/test_indexes.py (100%) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 0f0d0f309..95e184356 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -845,6 +845,8 @@ MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool] # Scratch types for generics +T_INT = TypeVar("T_INT", bound=int) +T_COMPLEX = TypeVar("T_COMPLEX", bound=complex) SeriesDType: TypeAlias = ( str | bytes diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index b193dccf8..03f7d686e 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -39,12 +39,15 @@ from pandas.core.strings.accessor import StringMethods from typing_extensions import ( Never, Self, + TypeAlias, ) from pandas._libs.interval import _OrderableT from pandas._typing import ( C2, S1, + T_COMPLEX, + T_INT, AnyAll, ArrayLike, AxesData, @@ -70,18 +73,32 @@ from pandas._typing import ( TimestampDtypeArg, np_1darray, np_ndarray_anyint, + np_ndarray_bool, np_ndarray_complex, np_ndarray_float, + np_ndarray_str, type_t, ) class InvalidIndexError(Exception): ... +_ListLike: TypeAlias = ArrayLike | dict[_str, np.ndarray] | SequenceNotStr[S1] + class Index(IndexOpsMixin[S1]): __hash__: ClassVar[None] # type: ignore[assignment] # overloads with additional dtypes @overload def __new__( # pyright: ignore[reportOverlappingOverload] + cls, + data: Sequence[bool | np.bool_] | IndexOpsMixin[bool] | np_ndarray_bool, + *, + dtype: Literal["bool"] | type_t[bool | np.bool_] = ..., + copy: bool = ..., + name: Hashable = ..., + tupleize_cols: bool = ..., + ) -> Index[bool]: ... + @overload + def __new__( cls, data: Sequence[int | np.integer] | IndexOpsMixin[int] | np_ndarray_anyint, *, @@ -460,6 +477,155 @@ class Index(IndexOpsMixin[S1]): def __gt__(self, other: Self | S1) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] # overwrite inherited methods from OpsMixin @overload + def __add__(self: Index[Never], other: _str) -> Never: ... + @overload + def __add__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ... + @overload + def __add__(self, other: Index[Never]) -> Index: ... + @overload + def __add__( + self: Index[bool], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX], + ) -> Index[T_COMPLEX]: ... + @overload + def __add__(self: Index[bool], other: np_ndarray_bool) -> Index[bool]: ... + @overload + def __add__(self: Index[bool], other: np_ndarray_anyint) -> Index[int]: ... + @overload + def __add__(self: Index[bool], other: np_ndarray_float) -> Index[float]: ... + @overload + def __add__(self: Index[bool], other: np_ndarray_complex) -> Index[complex]: ... + @overload + def __add__( + self: Index[int], + other: ( + bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Index[bool] + ), + ) -> Index[int]: ... + @overload + def __add__( + self: Index[int], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX], + ) -> Index[T_COMPLEX]: ... + @overload + def __add__(self: Index[int], other: np_ndarray_float) -> Index[float]: ... + @overload + def __add__(self: Index[int], other: np_ndarray_complex) -> Index[complex]: ... + @overload + def __add__( + self: Index[float], + other: ( + int + | Sequence[int] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | Index[T_INT] + ), + ) -> Index[float]: ... + @overload + def __add__( + self: Index[float], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX], + ) -> Index[T_COMPLEX]: ... + @overload + def __add__(self: Index[float], other: np_ndarray_complex) -> Index[complex]: ... + @overload + def __add__( + self: Index[complex], + other: ( + T_COMPLEX + | Sequence[T_COMPLEX] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + | Index[T_COMPLEX] + ), + ) -> Index[complex]: ... + @overload + def __add__( + self: Index[_str], + other: ( + np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex + ), + ) -> Never: ... + @overload + def __add__( + self: Index[_str], other: _str | Sequence[_str] | np_ndarray_str | Index[_str] + ) -> Index[_str]: ... + @overload # type: ignore[override] + def __radd__(self: Index[Never], other: _str) -> Never: ... + @overload + def __radd__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ... + @overload + def __radd__( + self: Index[bool], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX], + ) -> Index[T_COMPLEX]: ... + @overload + def __radd__(self: Index[bool], other: np_ndarray_bool) -> Index[bool]: ... + @overload + def __radd__(self: Index[bool], other: np_ndarray_anyint) -> Index[int]: ... + @overload + def __radd__(self: Index[bool], other: np_ndarray_float) -> Index[float]: ... + @overload + def __radd__( + self: Index[int], + other: ( + bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Index[bool] + ), + ) -> Index[int]: ... + @overload + def __radd__( + self: Index[int], other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] + ) -> Index[T_COMPLEX]: ... + @overload + def __radd__(self: Index[int], other: np_ndarray_float) -> Index[float]: ... + @overload + def __radd__( + self: Index[float], + other: ( + int + | Sequence[int] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | Index[T_INT] + ), + ) -> Index[float]: ... + @overload + def __radd__( + self: Index[float], other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] + ) -> Index[T_COMPLEX]: ... + @overload + def __radd__( + self: Index[complex], + other: ( + T_COMPLEX + | Sequence[T_COMPLEX] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | Index[T_COMPLEX] + ), + ) -> Index[complex]: ... + @overload + def __radd__( + self: Index[T_COMPLEX], other: np_ndarray_complex + ) -> Index[complex]: ... + @overload + def __radd__( + self: Index[_str], + other: ( + np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex + ), + ) -> Never: ... + @overload + def __radd__( + self: Index[_str], other: _str | Sequence[_str] | np_ndarray_str | Index[_str] + ) -> Index[_str]: ... + @overload def __mul__( self: Index[int] | Index[float], other: timedelta ) -> TimedeltaIndex: ... diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index 29c8f1123..3e4e14f61 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -61,10 +61,10 @@ class DatetimeIndex( def __reduce__(self): ... # various ignores needed for mypy, as we do want to restrict what can be used in # arithmetic for these types - @overload + @overload # type: ignore[override] def __add__(self, other: TimedeltaSeries) -> TimestampSeries: ... @overload - def __add__( + def __add__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: timedelta | Timedelta | TimedeltaIndex | BaseOffset ) -> DatetimeIndex: ... @overload diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index eb0d073c1..309770dba 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -49,12 +49,14 @@ class TimedeltaIndex( ) -> Self: ... # various ignores needed for mypy, as we do want to restrict what can be used in # arithmetic for these types - @overload + @overload # type: ignore[override] def __add__(self, other: Period) -> PeriodIndex: ... @overload def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ... @overload - def __add__(self, other: dt.timedelta | Timedelta | Self) -> Self: ... + def __add__( # pyright: ignore[reportIncompatibleMethodOverride] + self, other: dt.timedelta | Timedelta | Self + ) -> Self: ... def __radd__(self, other: dt.datetime | Timestamp | DatetimeIndex) -> DatetimeIndex: ... # type: ignore[override] def __sub__(self, other: dt.timedelta | Timedelta | Self) -> Self: ... def __mul__(self, other: num) -> Self: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 768e29b7c..0cc2320a9 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -26,7 +26,6 @@ from typing import ( Generic, Literal, NoReturn, - TypeVar, final, overload, type_check_only, @@ -108,6 +107,8 @@ from pandas._libs.tslibs.offsets import DateOffset from pandas._typing import ( S1, S2, + T_COMPLEX, + T_INT, AggFuncTypeBase, AggFuncTypeDictFrame, AggFuncTypeSeriesToFrame, @@ -199,9 +200,6 @@ from pandas.core.dtypes.dtypes import CategoricalDtype from pandas.plotting import PlotAccessor -_T_INT = TypeVar("_T_INT", bound=int) -_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex) - class _iLocIndexerSeries(_iLocIndexer, Generic[S1]): # get item @overload @@ -1625,8 +1623,8 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __add__( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __add__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... @overload @@ -1645,8 +1643,8 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __add__( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @overload @@ -1660,27 +1658,27 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __add__( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __add__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ... @overload def __add__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload @@ -1713,11 +1711,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def add( self: Series[bool], @@ -1763,11 +1761,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def add( self: Series[int], @@ -1793,7 +1791,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -1802,11 +1800,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def add( self: Series[float], @@ -1819,13 +1817,13 @@ class Series(IndexOpsMixin[S1], NDFrame): def add( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -1844,8 +1842,8 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __radd__( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __radd__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... @overload @@ -1861,8 +1859,8 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[int]: ... @overload def __radd__( - self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX] - ) -> Series[_T_COMPLEX]: ... + self: Series[int], other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX] + ) -> Series[T_COMPLEX]: ... @overload def __radd__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @overload @@ -1874,29 +1872,29 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __radd__( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __radd__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __radd__( - self: Series[_T_COMPLEX], other: np_ndarray_complex + self: Series[T_COMPLEX], other: np_ndarray_complex ) -> Series[complex]: ... @overload def __radd__( @@ -1920,11 +1918,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def radd( self: Series[bool], @@ -1962,11 +1960,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def radd( self: Series[int], @@ -1984,7 +1982,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -1993,21 +1991,21 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def radd( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -2015,7 +2013,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def radd( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: np_ndarray_complex, level: Level | None = None, fill_value: float | None = None, @@ -2057,8 +2055,8 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __mul__( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __mul__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... @overload @@ -2075,8 +2073,8 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __mul__( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __mul__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @overload @@ -2088,29 +2086,29 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __mul__( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __mul__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __mul__( - self: Series[_T_COMPLEX], other: np_ndarray_complex + self: Series[T_COMPLEX], other: np_ndarray_complex ) -> Series[complex]: ... @overload def __mul__( @@ -2135,11 +2133,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def mul( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def mul( self: Series[bool], @@ -2177,11 +2175,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def mul( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def mul( self: Series[int], @@ -2199,7 +2197,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -2208,21 +2206,21 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def mul( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def mul( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -2230,7 +2228,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def mul( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: np_ndarray_complex, level: Level | None = None, fill_value: float | None = None, @@ -2253,8 +2251,8 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __rmul__( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __rmul__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... @overload @@ -2270,8 +2268,8 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[int]: ... @overload def __rmul__( - self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX] - ) -> Series[_T_COMPLEX]: ... + self: Series[int], other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX] + ) -> Series[T_COMPLEX]: ... @overload def __rmul__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @overload @@ -2283,29 +2281,29 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __rmul__( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __rmul__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __rmul__( - self: Series[_T_COMPLEX], other: np_ndarray_complex + self: Series[T_COMPLEX], other: np_ndarray_complex ) -> Series[complex]: ... @overload def __rmul__( @@ -2330,11 +2328,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def rmul( self: Series[bool], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def rmul( self: Series[bool], @@ -2372,11 +2370,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def rmul( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def rmul( self: Series[int], @@ -2394,7 +2392,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -2403,21 +2401,21 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def rmul( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def rmul( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -2425,7 +2423,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def rmul( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: np_ndarray_complex, level: Level | None = None, fill_value: float | None = None, @@ -2528,17 +2526,17 @@ class Series(IndexOpsMixin[S1], NDFrame): def __sub__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __sub__( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: ( Just[complex] | Sequence[Just[complex]] @@ -2652,12 +2650,12 @@ class Series(IndexOpsMixin[S1], NDFrame): def sub( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -2665,7 +2663,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def sub( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: ( Just[complex] | Sequence[Just[complex]] @@ -2760,17 +2758,17 @@ class Series(IndexOpsMixin[S1], NDFrame): def __rsub__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __rsub__( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: ( Just[complex] | Sequence[Just[complex]] @@ -2862,12 +2860,12 @@ class Series(IndexOpsMixin[S1], NDFrame): def rsub( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -2875,7 +2873,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def rsub( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: ( Just[complex] | Sequence[Just[complex]] @@ -2921,14 +2919,14 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __truediv__( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __truediv__( self: Series[float], @@ -2938,29 +2936,29 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __truediv__( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __truediv__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __truediv__( - self: Series[_T_COMPLEX], other: np_ndarray_complex + self: Series[T_COMPLEX], other: np_ndarray_complex ) -> Series[complex]: ... @overload def __truediv__(self, other: Path) -> Series: ... @@ -3013,7 +3011,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -3022,11 +3020,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def truediv( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: AxisIndex = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def truediv( self: Series[float], @@ -3036,7 +3034,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -3045,21 +3043,21 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def truediv( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: AxisIndex = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def truediv( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -3067,7 +3065,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def truediv( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: np_ndarray_complex, level: Level | None = None, fill_value: float | None = None, @@ -3117,13 +3115,13 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __rtruediv__( - self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX] - ) -> Series[_T_COMPLEX]: ... + self: Series[int], other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX] + ) -> Series[T_COMPLEX]: ... @overload def __rtruediv__( self: Series[float], @@ -3133,29 +3131,29 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), ) -> Series[float]: ... @overload def __rtruediv__( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], - ) -> Series[_T_COMPLEX]: ... + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + ) -> Series[T_COMPLEX]: ... @overload def __rtruediv__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), ) -> Series[complex]: ... @overload def __rtruediv__( - self: Series[_T_COMPLEX], other: np_ndarray_complex + self: Series[T_COMPLEX], other: np_ndarray_complex ) -> Series[complex]: ... @overload def __rtruediv__(self, other: Path) -> Series: ... @@ -3208,7 +3206,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -3217,11 +3215,11 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def rtruediv( self: Series[int], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: AxisIndex = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def rtruediv( self: Series[float], @@ -3231,7 +3229,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_INT] + | Series[T_INT] ), level: Level | None = None, fill_value: float | None = None, @@ -3240,21 +3238,21 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def rtruediv( self: Series[float], - other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: AxisIndex = 0, - ) -> Series[_T_COMPLEX]: ... + ) -> Series[T_COMPLEX]: ... @overload def rtruediv( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] + T_COMPLEX + | Sequence[T_COMPLEX] | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float - | Series[_T_COMPLEX] + | Series[T_COMPLEX] ), level: Level | None = None, fill_value: float | None = None, @@ -3262,7 +3260,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[complex]: ... @overload def rtruediv( - self: Series[_T_COMPLEX], + self: Series[T_COMPLEX], other: np_ndarray_complex, level: Level | None = None, fill_value: float | None = None, diff --git a/tests/indexes/__init__.py b/tests/indexes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/__init__.py b/tests/indexes/arithmetic/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/bool/__init__.py b/tests/indexes/arithmetic/bool/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/bool/test_add.py b/tests/indexes/arithmetic/bool/test_add.py new file mode 100644 index 000000000..1f1392679 --- /dev/null +++ b/tests/indexes/arithmetic/bool/test_add.py @@ -0,0 +1,83 @@ +import numpy as np +from numpy import typing as npt # noqa: F401 +import pandas as pd +from typing_extensions import assert_type + +from tests import check + +# left operand +left = pd.Index([True, True, False]) + + +def test_add_py_scalar() -> None: + """Test pd.Index[bool] + Python native scalars""" + b, i, f, c = True, 1, 1.0, 1j + + left.__add__(b) + check(assert_type(left + b, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(i + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_py_sequence() -> None: + """Test pd.Index[bool] + Python native sequence""" + b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + b, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(i + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_numpy_array() -> None: + """Test pd.Index[bool] + numpy array""" + b = np.array([True, False, True], np.bool_) + i = np.array([2, 3, 5], np.int64) + f = np.array([1.0, 2.0, 3.0], np.float64) + c = np.array([1.1j, 2.2j, 4.1j], np.complex128) + + check(assert_type(left + b, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Index`s with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Index, np.bool_) + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Index, np.integer) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Index, np.floating) + check( + assert_type(c + left, "npt.NDArray[np.complex128]"), + pd.Index, + np.complexfloating, + ) + + +def test_add_pd_series() -> None: + """Test pd.Index[bool] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[bool]"), pd.Index, np.bool_) + check(assert_type(i + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) diff --git a/tests/indexes/arithmetic/complex/__init__.py b/tests/indexes/arithmetic/complex/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/complex/test_add.py b/tests/indexes/arithmetic/complex/test_add.py new file mode 100644 index 000000000..9dda8fe2b --- /dev/null +++ b/tests/indexes/arithmetic/complex/test_add.py @@ -0,0 +1,84 @@ +import numpy as np +from numpy import typing as npt # noqa: F401 +import pandas as pd +from typing_extensions import assert_type + +from tests import check + +# left operand +left = pd.Index([1j, 2j, 3j]) + + +def test_add_py_scalar() -> None: + """Test pd.Index[complex] + Python native scalars""" + b, i, f, c = True, 1, 1.0, 1j + + check(assert_type(left + b, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + i, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + f, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(i + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(f + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_py_sequence() -> None: + """Test pd.Index[complex] + Python native sequence""" + b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + b, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + i, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + f, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(i + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(f + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_numpy_array() -> None: + """Test pd.Index[complex] + numpy array""" + b = np.array([True, False, True], np.bool_) + i = np.array([2, 3, 5], np.int64) + f = np.array([1.0, 2.0, 3.0], np.float64) + c = np.array([1.1j, 2.2j, 4.1j], np.complex128) + + check(assert_type(left + b, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + i, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + f, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Index`s with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Index, np.complexfloating) + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Index, np.complexfloating) + check( + assert_type(f + left, "npt.NDArray[np.float64]"), pd.Index, np.complexfloating + ) + check( + assert_type(c + left, "npt.NDArray[np.complex128]"), + pd.Index, + np.complexfloating, + ) + + +def test_add_pd_series() -> None: + """Test pd.Index[complex] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + i, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + f, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(i + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(f + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) diff --git a/tests/indexes/arithmetic/float/__init__.py b/tests/indexes/arithmetic/float/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/float/test_add.py b/tests/indexes/arithmetic/float/test_add.py new file mode 100644 index 000000000..26f6ba964 --- /dev/null +++ b/tests/indexes/arithmetic/float/test_add.py @@ -0,0 +1,82 @@ +import numpy as np +from numpy import typing as npt # noqa: F401 +import pandas as pd +from typing_extensions import assert_type + +from tests import check + +# left operand +left = pd.Index([1.0, 2.0, 3.0]) + + +def test_add_py_scalar() -> None: + """Test pd.Index[float] + Python native scalars""" + b, i, f, c = True, 1, 1.0, 1j + + check(assert_type(left + b, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + i, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(i + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_py_sequence() -> None: + """Test pd.Index[float] + Python native sequence""" + b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + b, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + i, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(i + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_numpy_array() -> None: + """Test pd.Index[float] + numpy array""" + b = np.array([True, False, True], np.bool_) + i = np.array([2, 3, 5], np.int64) + f = np.array([1.0, 2.0, 3.0], np.float64) + c = np.array([1.1j, 2.2j, 4.1j], np.complex128) + + check(assert_type(left + b, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + i, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Index`s with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Index, np.floating) + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Index, np.floating) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Index, np.floating) + check( + assert_type(c + left, "npt.NDArray[np.complex128]"), + pd.Index, + np.complexfloating, + ) + + +def test_add_pd_series() -> None: + """Test pd.Index[float] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + i, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(i + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) diff --git a/tests/indexes/arithmetic/int/__init__.py b/tests/indexes/arithmetic/int/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/int/test_add.py b/tests/indexes/arithmetic/int/test_add.py new file mode 100644 index 000000000..01eebd828 --- /dev/null +++ b/tests/indexes/arithmetic/int/test_add.py @@ -0,0 +1,82 @@ +import numpy as np +from numpy import typing as npt # noqa: F401 +import pandas as pd +from typing_extensions import assert_type + +from tests import check + +# left operand +left = pd.Index([1, 2, 3]) + + +def test_add_py_scalar() -> None: + """Test pd.Index[int] + Python native scalars""" + b, i, f, c = True, 1, 1.0, 1j + + check(assert_type(left + b, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_py_sequence() -> None: + """Test pd.Index[int] + Python native sequence""" + b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + b, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) + + +def test_add_numpy_array() -> None: + """Test pd.Index[int] + numpy array""" + b = np.array([True, False, True], np.bool_) + i = np.array([2, 3, 5], np.int64) + f = np.array([1.0, 2.0, 3.0], np.float64) + c = np.array([1.1j, 2.2j, 4.1j], np.complex128) + + check(assert_type(left + b, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Index`s with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Index, np.integer) + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Index, np.integer) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Index, np.floating) + check( + assert_type(c + left, "npt.NDArray[np.complex128]"), + pd.Index, + np.complexfloating, + ) + + +def test_add_pd_series() -> None: + """Test pd.Index[int] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + i, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(left + f, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(left + c, "pd.Index[complex]"), pd.Index, np.complexfloating) + + check(assert_type(b + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i + left, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(f + left, "pd.Index[float]"), pd.Index, np.floating) + check(assert_type(c + left, "pd.Index[complex]"), pd.Index, np.complexfloating) diff --git a/tests/indexes/arithmetic/str/__init__.py b/tests/indexes/arithmetic/str/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/indexes/arithmetic/str/test_add.py b/tests/indexes/arithmetic/str/test_add.py new file mode 100644 index 000000000..b48dc844c --- /dev/null +++ b/tests/indexes/arithmetic/str/test_add.py @@ -0,0 +1,88 @@ +import sys +from typing import Any + +import numpy as np +from numpy import typing as npt # noqa: F401 +import pandas as pd +from typing_extensions import ( + Never, + assert_type, +) + +from tests import ( + TYPE_CHECKING_INVALID_USAGE, + check, +) + +left = pd.Index(["1", "23", "456"]) # left operand + + +def test_add_py_scalar() -> None: + """Test pd.Index[str] + Python native 'scalar's""" + i = 4 + r0 = "right" + + if TYPE_CHECKING_INVALID_USAGE: + _0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(left + r0, "pd.Index[str]"), pd.Index, str) + + if TYPE_CHECKING_INVALID_USAGE: + _1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(r0 + left, "pd.Index[str]"), pd.Index, str) + + +def test_add_py_sequence() -> None: + """Test pd.Index[str] + Python native sequence""" + i = [3, 5, 8] + r0 = ["a", "bc", "def"] + r1 = tuple(r0) + + if TYPE_CHECKING_INVALID_USAGE: + _0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(left + r0, "pd.Index[str]"), pd.Index, str) + check(assert_type(left + r1, "pd.Index[str]"), pd.Index, str) + + if TYPE_CHECKING_INVALID_USAGE: + _1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(r0 + left, "pd.Index[str]"), pd.Index, str) + check(assert_type(r1 + left, "pd.Index[str]"), pd.Index, str) + + +def test_add_numpy_array() -> None: + """Test pd.Index[str] + numpy array""" + i = np.array([3, 5, 8], np.int64) + r0 = np.array(["a", "bc", "def"], np.str_) + + if TYPE_CHECKING_INVALID_USAGE: + assert_type(left + i, Never) + check(assert_type(left + r0, "pd.Index[str]"), pd.Index, str) + + # `numpy` typing gives `npt.NDArray[np.int64]` in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Index`s. + if TYPE_CHECKING_INVALID_USAGE: + assert_type(i + left, "npt.NDArray[np.int64]") + if sys.version_info >= (3, 11): + # `numpy` typing gives `npt.NDArray[np.int64]` in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Index`s. + check(assert_type(r0 + left, "npt.NDArray[np.str_]"), pd.Index, str) + else: + # Python 3.10 uses NumPy 2.2.6, and it has for r0 ndarray[tuple[int,...], dtype[str_]] + # Python 3.11+ uses NumPy 2.3.2, and it has for r0 ndarray[tuple[Any,...,dtype[str_]] + # https://github.com/pandas-dev/pandas-stubs/pull/1274#discussion_r2291498975 + check(assert_type(r0 + left, Any), pd.Index, str) + + +def test_add_pd_series() -> None: + """Test pd.Index[str] + pandas index""" + i = pd.Index([3, 5, 8]) + r0 = pd.Index(["a", "bc", "def"]) + + if TYPE_CHECKING_INVALID_USAGE: + _0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(left + r0, "pd.Index[str]"), pd.Index, str) + + if TYPE_CHECKING_INVALID_USAGE: + _1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(r0 + left, "pd.Index[str]"), pd.Index, str) diff --git a/tests/indexes/arithmetic/test_add.py b/tests/indexes/arithmetic/test_add.py new file mode 100644 index 000000000..bc3c75eaa --- /dev/null +++ b/tests/indexes/arithmetic/test_add.py @@ -0,0 +1,87 @@ +import numpy as np +from numpy import typing as npt # noqa: F401 +import pandas as pd +from typing_extensions import assert_type + +from tests import check + +# left operand +left_i = pd.MultiIndex.from_tuples([(1,), (2,), (3,)]).levels[0] + + +def test_add_py_scalar() -> None: + """Test pd.Index[Any] + Python native scalars""" + b, i, f, c = True, 1, 1.0, 1j + + check(assert_type(left_i + b, pd.Index), pd.Index) + check(assert_type(left_i + i, pd.Index), pd.Index) + check(assert_type(left_i + f, pd.Index), pd.Index) + check(assert_type(left_i + c, pd.Index), pd.Index) + + check(assert_type(b + left_i, pd.Index), pd.Index) + check(assert_type(i + left_i, pd.Index), pd.Index) + check(assert_type(f + left_i, pd.Index), pd.Index) + check(assert_type(c + left_i, pd.Index), pd.Index) + + +def test_add_py_sequence() -> None: + """Test pd.Index[Any] + Python native sequence""" + b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left_i + b, pd.Index), pd.Index) + check(assert_type(left_i + i, pd.Index), pd.Index) + check(assert_type(left_i + f, pd.Index), pd.Index) + check(assert_type(left_i + c, pd.Index), pd.Index) + + check(assert_type(b + left_i, pd.Index), pd.Index) + check(assert_type(i + left_i, pd.Index), pd.Index) + check(assert_type(f + left_i, pd.Index), pd.Index) + check(assert_type(c + left_i, pd.Index), pd.Index) + + +def test_add_numpy_array() -> None: + """Test pd.Index[Any] + numpy array""" + b = np.array([True, False, True], np.bool_) + i = np.array([2, 3, 5], np.int64) + f = np.array([1.0, 2.0, 3.0], np.float64) + c = np.array([1.1j, 2.2j, 4.1j], np.complex128) + + check(assert_type(left_i + b, pd.Index), pd.Index) + check(assert_type(left_i + i, pd.Index), pd.Index) + check(assert_type(left_i + f, pd.Index), pd.Index) + check(assert_type(left_i + c, pd.Index), pd.Index) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Series`s. + # `mypy` thinks the return types are `Any`, which is a bug. + check( + assert_type(b + left_i, "npt.NDArray[np.bool_]"), pd.Index # type: ignore[assert-type] + ) + check( + assert_type(i + left_i, "npt.NDArray[np.int64]"), pd.Index # type: ignore[assert-type] + ) + check( + assert_type(f + left_i, "npt.NDArray[np.float64]"), pd.Index # type: ignore[assert-type] + ) + check( + assert_type(c + left_i, "npt.NDArray[np.complex128]"), pd.Index # type: ignore[assert-type] + ) + + +def test_add_pd_series() -> None: + """Test pd.Index[Any] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left_i + b, pd.Index), pd.Index) + check(assert_type(left_i + i, pd.Index), pd.Index) + check(assert_type(left_i + f, pd.Index), pd.Index) + check(assert_type(left_i + c, pd.Index), pd.Index) + + check(assert_type(b + left_i, pd.Index), pd.Index) + check(assert_type(i + left_i, pd.Index), pd.Index) + check(assert_type(f + left_i, pd.Index), pd.Index) + check(assert_type(c + left_i, pd.Index), pd.Index) diff --git a/tests/test_indexes.py b/tests/indexes/test_indexes.py similarity index 100% rename from tests/test_indexes.py rename to tests/indexes/test_indexes.py From 9f7b696a8d3c91953738f2df87d1eadab97cd119 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Wed, 27 Aug 2025 22:56:49 +0200 Subject: [PATCH 2/4] chore: missing tests for Index[Any] --- tests/indexes/arithmetic/test_add.py | 43 ++++++++++++++++++++-------- 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/tests/indexes/arithmetic/test_add.py b/tests/indexes/arithmetic/test_add.py index bc3c75eaa..b82a51c02 100644 --- a/tests/indexes/arithmetic/test_add.py +++ b/tests/indexes/arithmetic/test_add.py @@ -1,16 +1,23 @@ import numpy as np from numpy import typing as npt # noqa: F401 import pandas as pd -from typing_extensions import assert_type +from typing_extensions import ( + Never, + assert_type, +) -from tests import check +from tests import ( + TYPE_CHECKING_INVALID_USAGE, + check, +) -# left operand +# left operands left_i = pd.MultiIndex.from_tuples([(1,), (2,), (3,)]).levels[0] +left_str = pd.MultiIndex.from_tuples([("1",), ("2",), ("3_",)]).levels[0] -def test_add_py_scalar() -> None: - """Test pd.Index[Any] + Python native scalars""" +def test_add_i_py_scalar() -> None: + """Test pd.Index[Any] (int) + Python native scalars""" b, i, f, c = True, 1, 1.0, 1j check(assert_type(left_i + b, pd.Index), pd.Index) @@ -24,8 +31,8 @@ def test_add_py_scalar() -> None: check(assert_type(c + left_i, pd.Index), pd.Index) -def test_add_py_sequence() -> None: - """Test pd.Index[Any] + Python native sequence""" +def test_add_i_py_sequence() -> None: + """Test pd.Index[Any] (int) + Python native sequence""" b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] check(assert_type(left_i + b, pd.Index), pd.Index) @@ -39,8 +46,8 @@ def test_add_py_sequence() -> None: check(assert_type(c + left_i, pd.Index), pd.Index) -def test_add_numpy_array() -> None: - """Test pd.Index[Any] + numpy array""" +def test_add_i_numpy_array() -> None: + """Test pd.Index[Any] (int) + numpy array""" b = np.array([True, False, True], np.bool_) i = np.array([2, 3, 5], np.int64) f = np.array([1.0, 2.0, 3.0], np.float64) @@ -53,7 +60,7 @@ def test_add_numpy_array() -> None: # `numpy` typing gives the corresponding `ndarray`s in the static type # checking, where our `__radd__` cannot override. At runtime, they return - # `Series`s. + # `Index`s. # `mypy` thinks the return types are `Any`, which is a bug. check( assert_type(b + left_i, "npt.NDArray[np.bool_]"), pd.Index # type: ignore[assert-type] @@ -69,19 +76,31 @@ def test_add_numpy_array() -> None: ) -def test_add_pd_series() -> None: - """Test pd.Index[Any] + pandas index""" +def test_add_i_pd_series() -> None: + """Test pd.Index[Any] (int) + pandas index""" + a = pd.MultiIndex.from_tuples([(1,), (2,), (3,)]).levels[0] b = pd.Index([True, False, True]) i = pd.Index([2, 3, 5]) f = pd.Index([1.0, 2.0, 3.0]) c = pd.Index([1.1j, 2.2j, 4.1j]) + check(assert_type(left_i + a, pd.Index), pd.Index) check(assert_type(left_i + b, pd.Index), pd.Index) check(assert_type(left_i + i, pd.Index), pd.Index) check(assert_type(left_i + f, pd.Index), pd.Index) check(assert_type(left_i + c, pd.Index), pd.Index) + check(assert_type(a + left_i, pd.Index), pd.Index) check(assert_type(b + left_i, pd.Index), pd.Index) check(assert_type(i + left_i, pd.Index), pd.Index) check(assert_type(f + left_i, pd.Index), pd.Index) check(assert_type(c + left_i, pd.Index), pd.Index) + + +def test_add_i_py_str() -> None: + """Test pd.Index[Any] (int) + Python str""" + s = "abc" + + if TYPE_CHECKING_INVALID_USAGE: + assert_type(left_i + s, Never) + assert_type(s + left_i, Never) From 1a1ff50893585e0900d3906508210f25949495b4 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Thu, 28 Aug 2025 23:55:41 +0200 Subject: [PATCH 3/4] feat(comment): https://github.com/pandas-dev/pandas-stubs/pull/1350#discussion_r2305349459 --- pandas-stubs/core/series.pyi | 83 ++++++++++++++------ tests/indexes/arithmetic/bool/test_add.py | 2 +- tests/indexes/arithmetic/complex/test_add.py | 2 +- tests/indexes/arithmetic/float/test_add.py | 2 +- tests/indexes/arithmetic/int/test_add.py | 2 +- tests/indexes/arithmetic/str/test_add.py | 2 +- tests/indexes/arithmetic/test_add.py | 2 +- tests/series/arithmetic/bool/test_add.py | 30 +++++++ tests/series/arithmetic/complex/test_add.py | 36 +++++++++ tests/series/arithmetic/float/test_add.py | 30 +++++++ tests/series/arithmetic/int/test_add.py | 30 +++++++ tests/series/arithmetic/str/test_add.py | 22 ++++++ tests/series/arithmetic/test_add.py | 49 +++++++++++- 13 files changed, 260 insertions(+), 32 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 1b823cdc0..ef3c625ff 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1631,13 +1631,15 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __add__(self: Series[Never], other: _str) -> Never: ... @overload - def __add__(self: Series[Never], other: complex | _ListLike | Series) -> Series: ... + def __add__( + self: Series[Never], other: complex | _ListLike | Index | Series + ) -> Series: ... @overload - def __add__(self, other: Series[Never]) -> Series: ... + def __add__(self, other: Index[Never] | Series[Never]) -> Series: ... @overload def __add__( self: Series[bool], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], ) -> Series[T_COMPLEX]: ... @overload def __add__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... @@ -1651,13 +1653,18 @@ class Series(IndexOpsMixin[S1], NDFrame): def __add__( self: Series[int], other: ( - bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + bool + | Sequence[bool] + | np_ndarray_bool + | np_ndarray_anyint + | Index[bool] + | Series[bool] ), ) -> Series[int]: ... @overload def __add__( self: Series[int], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], ) -> Series[T_COMPLEX]: ... @overload def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @@ -1672,13 +1679,14 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float + | Index[T_INT] | Series[T_INT] ), ) -> Series[float]: ... @overload def __add__( self: Series[float], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], ) -> Series[T_COMPLEX]: ... @overload def __add__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ... @@ -1692,6 +1700,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex + | Index[T_COMPLEX] | Series[T_COMPLEX] ), ) -> Series[complex]: ... @@ -1704,7 +1713,8 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Never: ... @overload def __add__( - self: Series[_str], other: _str | Sequence[_str] | np_ndarray_str | Series[_str] + self: Series[_str], + other: _str | Sequence[_str] | np_ndarray_str | Index[_str] | Series[_str], ) -> Series[_str]: ... @overload def add( @@ -1717,7 +1727,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[Never], - other: complex | _ListLike | Series, + other: complex | _ListLike | Index | Series, level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1733,7 +1743,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[bool], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1774,7 +1784,12 @@ class Series(IndexOpsMixin[S1], NDFrame): def add( self: Series[int], other: ( - bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + bool + | Sequence[bool] + | np_ndarray_bool + | np_ndarray_anyint + | Index[bool] + | Series[bool] ), level: Level | None = None, fill_value: float | None = None, @@ -1783,7 +1798,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[int], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1813,6 +1828,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float + | Index[T_INT] | Series[T_INT] ), level: Level | None = None, @@ -1822,7 +1838,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[float], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1845,6 +1861,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex + | Index[T_COMPLEX] | Series[T_COMPLEX] ), level: Level | None = None, @@ -1854,7 +1871,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[_str], - other: _str | Sequence[_str] | np_ndarray_str | Series[_str], + other: _str | Sequence[_str] | np_ndarray_str | Index[_str] | Series[_str], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1863,12 +1880,12 @@ class Series(IndexOpsMixin[S1], NDFrame): def __radd__(self: Series[Never], other: _str) -> Never: ... @overload def __radd__( - self: Series[Never], other: complex | _ListLike | Series + self: Series[Never], other: complex | _ListLike | Index | Series ) -> Series: ... @overload def __radd__( self: Series[bool], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], ) -> Series[T_COMPLEX]: ... @overload def __radd__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... @@ -1880,12 +1897,18 @@ class Series(IndexOpsMixin[S1], NDFrame): def __radd__( self: Series[int], other: ( - bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + bool + | Sequence[bool] + | np_ndarray_bool + | np_ndarray_anyint + | Index[bool] + | Series[bool] ), ) -> Series[int]: ... @overload def __radd__( - self: Series[int], other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX] + self: Series[int], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], ) -> Series[T_COMPLEX]: ... @overload def __radd__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @@ -1898,13 +1921,14 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float + | Index[T_INT] | Series[T_INT] ), ) -> Series[float]: ... @overload def __radd__( self: Series[float], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], ) -> Series[T_COMPLEX]: ... @overload def __radd__( @@ -1915,6 +1939,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float + | Index[T_COMPLEX] | Series[T_COMPLEX] ), ) -> Series[complex]: ... @@ -1931,7 +1956,8 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Never: ... @overload def __radd__( - self: Series[_str], other: _str | Sequence[_str] | np_ndarray_str | Series[_str] + self: Series[_str], + other: _str | Sequence[_str] | np_ndarray_str | Index[_str] | Series[_str], ) -> Series[_str]: ... @overload def radd( @@ -1944,7 +1970,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[Never], - other: complex | _ListLike | Series, + other: complex | _ListLike | Index | Series, level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1960,7 +1986,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[bool], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1993,7 +2019,12 @@ class Series(IndexOpsMixin[S1], NDFrame): def radd( self: Series[int], other: ( - bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + bool + | Sequence[bool] + | np_ndarray_bool + | np_ndarray_anyint + | Index[bool] + | Series[bool] ), level: Level | None = None, fill_value: float | None = None, @@ -2002,7 +2033,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[int], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -2024,6 +2055,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float + | Index[T_INT] | Series[T_INT] ), level: Level | None = None, @@ -2033,7 +2065,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[float], - other: T_COMPLEX | Sequence[T_COMPLEX] | Series[T_COMPLEX], + other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX] | Series[T_COMPLEX], level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -2047,6 +2079,7 @@ class Series(IndexOpsMixin[S1], NDFrame): | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float + | Index[T_COMPLEX] | Series[T_COMPLEX] ), level: Level | None = None, @@ -2064,7 +2097,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[_str], - other: _str | Sequence[_str] | np_ndarray_str | Series[_str], + other: _str | Sequence[_str] | np_ndarray_str | Index[_str] | Series[_str], level: Level | None = None, fill_value: float | None = None, axis: int = 0, diff --git a/tests/indexes/arithmetic/bool/test_add.py b/tests/indexes/arithmetic/bool/test_add.py index 1f1392679..1d826a902 100644 --- a/tests/indexes/arithmetic/bool/test_add.py +++ b/tests/indexes/arithmetic/bool/test_add.py @@ -65,7 +65,7 @@ def test_add_numpy_array() -> None: ) -def test_add_pd_series() -> None: +def test_add_pd_index() -> None: """Test pd.Index[bool] + pandas index""" b = pd.Index([True, False, True]) i = pd.Index([2, 3, 5]) diff --git a/tests/indexes/arithmetic/complex/test_add.py b/tests/indexes/arithmetic/complex/test_add.py index 9dda8fe2b..4454664f6 100644 --- a/tests/indexes/arithmetic/complex/test_add.py +++ b/tests/indexes/arithmetic/complex/test_add.py @@ -66,7 +66,7 @@ def test_add_numpy_array() -> None: ) -def test_add_pd_series() -> None: +def test_add_pd_index() -> None: """Test pd.Index[complex] + pandas index""" b = pd.Index([True, False, True]) i = pd.Index([2, 3, 5]) diff --git a/tests/indexes/arithmetic/float/test_add.py b/tests/indexes/arithmetic/float/test_add.py index 26f6ba964..12809c17b 100644 --- a/tests/indexes/arithmetic/float/test_add.py +++ b/tests/indexes/arithmetic/float/test_add.py @@ -64,7 +64,7 @@ def test_add_numpy_array() -> None: ) -def test_add_pd_series() -> None: +def test_add_pd_index() -> None: """Test pd.Index[float] + pandas index""" b = pd.Index([True, False, True]) i = pd.Index([2, 3, 5]) diff --git a/tests/indexes/arithmetic/int/test_add.py b/tests/indexes/arithmetic/int/test_add.py index 01eebd828..8c714e88f 100644 --- a/tests/indexes/arithmetic/int/test_add.py +++ b/tests/indexes/arithmetic/int/test_add.py @@ -64,7 +64,7 @@ def test_add_numpy_array() -> None: ) -def test_add_pd_series() -> None: +def test_add_pd_index() -> None: """Test pd.Index[int] + pandas index""" b = pd.Index([True, False, True]) i = pd.Index([2, 3, 5]) diff --git a/tests/indexes/arithmetic/str/test_add.py b/tests/indexes/arithmetic/str/test_add.py index b48dc844c..1b8de718c 100644 --- a/tests/indexes/arithmetic/str/test_add.py +++ b/tests/indexes/arithmetic/str/test_add.py @@ -74,7 +74,7 @@ def test_add_numpy_array() -> None: check(assert_type(r0 + left, Any), pd.Index, str) -def test_add_pd_series() -> None: +def test_add_pd_index() -> None: """Test pd.Index[str] + pandas index""" i = pd.Index([3, 5, 8]) r0 = pd.Index(["a", "bc", "def"]) diff --git a/tests/indexes/arithmetic/test_add.py b/tests/indexes/arithmetic/test_add.py index b82a51c02..dea1ec7db 100644 --- a/tests/indexes/arithmetic/test_add.py +++ b/tests/indexes/arithmetic/test_add.py @@ -76,7 +76,7 @@ def test_add_i_numpy_array() -> None: ) -def test_add_i_pd_series() -> None: +def test_add_i_pd_index() -> None: """Test pd.Index[Any] (int) + pandas index""" a = pd.MultiIndex.from_tuples([(1,), (2,), (3,)]).levels[0] b = pd.Index([True, False, True]) diff --git a/tests/series/arithmetic/bool/test_add.py b/tests/series/arithmetic/bool/test_add.py index d2a77d55f..7ba110fc6 100644 --- a/tests/series/arithmetic/bool/test_add.py +++ b/tests/series/arithmetic/bool/test_add.py @@ -127,3 +127,33 @@ def test_add_pd_series() -> None: check( assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating ) + + +def test_add_pd_index() -> None: + """Test pd.Series[bool] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(b + left, "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.add(b), "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.radd(b), "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.floating) + check( + assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating + ) diff --git a/tests/series/arithmetic/complex/test_add.py b/tests/series/arithmetic/complex/test_add.py index e3668b338..5614a627b 100644 --- a/tests/series/arithmetic/complex/test_add.py +++ b/tests/series/arithmetic/complex/test_add.py @@ -153,3 +153,39 @@ def test_add_pd_series() -> None: check( assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating ) + + +def test_add_pd_index() -> None: + """Test pd.Series[complex] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(b + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.add(b), "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complexfloating) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complexfloating) + + check( + assert_type(left.radd(b), "pd.Series[complex]"), pd.Series, np.complexfloating + ) + check( + assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complexfloating + ) + check( + assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complexfloating + ) + check( + assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating + ) diff --git a/tests/series/arithmetic/float/test_add.py b/tests/series/arithmetic/float/test_add.py index 41fc6b96b..63ef7d907 100644 --- a/tests/series/arithmetic/float/test_add.py +++ b/tests/series/arithmetic/float/test_add.py @@ -127,3 +127,33 @@ def test_add_pd_series() -> None: check( assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating ) + + +def test_add_pd_index() -> None: + """Test pd.Series[float] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(b + left, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.add(b), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.radd(b), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.floating) + check( + assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating + ) diff --git a/tests/series/arithmetic/int/test_add.py b/tests/series/arithmetic/int/test_add.py index b68b5e235..cc2d7bd2e 100644 --- a/tests/series/arithmetic/int/test_add.py +++ b/tests/series/arithmetic/int/test_add.py @@ -127,3 +127,33 @@ def test_add_pd_series() -> None: check( assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating ) + + +def test_add_pd_index() -> None: + """Test pd.Series[int] + pandas index""" + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + b, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(b + left, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.add(b), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.floating) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complexfloating) + + check(assert_type(left.radd(b), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.floating) + check( + assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating + ) diff --git a/tests/series/arithmetic/str/test_add.py b/tests/series/arithmetic/str/test_add.py index a58fe663a..9668bdba7 100644 --- a/tests/series/arithmetic/str/test_add.py +++ b/tests/series/arithmetic/str/test_add.py @@ -120,3 +120,25 @@ def test_add_pd_series() -> None: if TYPE_CHECKING_INVALID_USAGE: left.radd(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue] check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str) + + +def test_add_pd_index() -> None: + """Test pd.Series[str] + pandas index""" + i = pd.Index([3, 5, 8]) + r0 = pd.Index(["a", "bc", "def"]) + + if TYPE_CHECKING_INVALID_USAGE: + _0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str) + + if TYPE_CHECKING_INVALID_USAGE: + _1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue] + check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str) + + if TYPE_CHECKING_INVALID_USAGE: + left.add(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue] + check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str) + + if TYPE_CHECKING_INVALID_USAGE: + left.radd(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue] + check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str) diff --git a/tests/series/arithmetic/test_add.py b/tests/series/arithmetic/test_add.py index a739ddae3..b0f867b59 100644 --- a/tests/series/arithmetic/test_add.py +++ b/tests/series/arithmetic/test_add.py @@ -1,3 +1,8 @@ +from typing import ( + Any, + cast, +) + import numpy as np from numpy import typing as npt # noqa: F401 import pandas as pd @@ -43,23 +48,32 @@ def test_add_i_py_scalar() -> None: def test_add_i_py_sequence() -> None: """Test pd.Series[Any] (int) + Python native sequence""" - b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + # mypy believe it is already list[Any], whereas pyright gives list[bool | int | complex] + a = cast(list[Any], [True, 3, 5j]) # type: ignore[redundant-cast] + b = [True, False, True] + i = [2, 3, 5] + f = [1.0, 2.0, 3.0] + c = [1j, 1j, 4j] + check(assert_type(left_i + a, pd.Series), pd.Series) check(assert_type(left_i + b, pd.Series), pd.Series) check(assert_type(left_i + i, pd.Series), pd.Series) check(assert_type(left_i + f, pd.Series), pd.Series) check(assert_type(left_i + c, pd.Series), pd.Series) + check(assert_type(a + left_i, pd.Series), pd.Series) check(assert_type(b + left_i, pd.Series), pd.Series) check(assert_type(i + left_i, pd.Series), pd.Series) check(assert_type(f + left_i, pd.Series), pd.Series) check(assert_type(c + left_i, pd.Series), pd.Series) + check(assert_type(left_i.add(a), pd.Series), pd.Series) check(assert_type(left_i.add(b), pd.Series), pd.Series) check(assert_type(left_i.add(i), pd.Series), pd.Series) check(assert_type(left_i.add(f), pd.Series), pd.Series) check(assert_type(left_i.add(c), pd.Series), pd.Series) + check(assert_type(left_i.radd(a), pd.Series), pd.Series) check(assert_type(left_i.radd(b), pd.Series), pd.Series) check(assert_type(left_i.radd(i), pd.Series), pd.Series) check(assert_type(left_i.radd(f), pd.Series), pd.Series) @@ -139,6 +153,39 @@ def test_add_i_pd_series() -> None: check(assert_type(left_i.radd(c), pd.Series), pd.Series) +def test_add_i_pd_index() -> None: + """Test pd.Series[Any] (int) + pandas index""" + a = pd.MultiIndex.from_tuples([(1,), (2,), (3,)]).levels[0] + b = pd.Index([True, False, True]) + i = pd.Index([2, 3, 5]) + f = pd.Index([1.0, 2.0, 3.0]) + c = pd.Index([1.1j, 2.2j, 4.1j]) + + check(assert_type(left_i + a, pd.Series), pd.Series) + check(assert_type(left_i + b, pd.Series), pd.Series) + check(assert_type(left_i + i, pd.Series), pd.Series) + check(assert_type(left_i + f, pd.Series), pd.Series) + check(assert_type(left_i + c, pd.Series), pd.Series) + + check(assert_type(a + left_i, pd.Series), pd.Series) + check(assert_type(b + left_i, pd.Series), pd.Series) + check(assert_type(i + left_i, pd.Series), pd.Series) + check(assert_type(f + left_i, pd.Series), pd.Series) + check(assert_type(c + left_i, pd.Series), pd.Series) + + check(assert_type(left_i.add(a), pd.Series), pd.Series) + check(assert_type(left_i.add(b), pd.Series), pd.Series) + check(assert_type(left_i.add(i), pd.Series), pd.Series) + check(assert_type(left_i.add(f), pd.Series), pd.Series) + check(assert_type(left_i.add(c), pd.Series), pd.Series) + + check(assert_type(left_i.radd(a), pd.Series), pd.Series) + check(assert_type(left_i.radd(b), pd.Series), pd.Series) + check(assert_type(left_i.radd(i), pd.Series), pd.Series) + check(assert_type(left_i.radd(f), pd.Series), pd.Series) + check(assert_type(left_i.radd(c), pd.Series), pd.Series) + + def test_add_i_py_str() -> None: """Test pd.Series[Any] (int) + Python str""" s = "abc" From 8a0dcbdb1fa01316c9c7b0493cfad9b27911dda8 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Fri, 29 Aug 2025 02:53:57 +0200 Subject: [PATCH 4/4] fix(comment): https://github.com/pandas-dev/pandas-stubs/pull/1350#pullrequestreview-3166747330 --- pandas-stubs/core/arraylike.pyi | 2 -- pandas-stubs/core/frame.pyi | 16 +++++++++------- pandas-stubs/core/indexes/base.pyi | 2 +- pandas-stubs/core/series.pyi | 2 +- tests/indexes/test_indexes.py | 2 -- 5 files changed, 11 insertions(+), 13 deletions(-) diff --git a/pandas-stubs/core/arraylike.pyi b/pandas-stubs/core/arraylike.pyi index 8eaaf5e13..ac981e291 100644 --- a/pandas-stubs/core/arraylike.pyi +++ b/pandas-stubs/core/arraylike.pyi @@ -19,8 +19,6 @@ class OpsMixin: def __rxor__(self, other: Any) -> Self: ... # ------------------------------------------------------------- # Arithmetic Methods - def __add__(self, other: Any) -> Self: ... - def __radd__(self, other: Any) -> Self: ... def __sub__(self, other: Any) -> Self: ... def __rsub__(self, other: Any) -> Self: ... def __mul__(self, other: Any) -> Self: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 10bff7b28..ae4c6a830 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1772,6 +1772,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): # methods @final def abs(self) -> Self: ... + def __add__(self, other: Any) -> Self: ... def add( self, other: num | ListLike | DataFrame, @@ -1779,6 +1780,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): level: Level | None = None, fill_value: float | None = None, ) -> Self: ... + def __radd__(self, other: Any) -> Self: ... + def radd( + self, + other, + axis: Axis = "columns", + level: Level | None = None, + fill_value: float | None = None, + ) -> Self: ... @final def add_prefix(self, prefix: _str, axis: Axis | None = None) -> Self: ... @final @@ -2222,13 +2231,6 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): min_count: int = 0, **kwargs: Any, ) -> Series: ... - def radd( - self, - other, - axis: Axis = "columns", - level: Level | None = None, - fill_value: float | None = None, - ) -> Self: ... @final def rank( self, diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 03f7d686e..14a02923c 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -554,7 +554,7 @@ class Index(IndexOpsMixin[S1]): def __add__( self: Index[_str], other: _str | Sequence[_str] | np_ndarray_str | Index[_str] ) -> Index[_str]: ... - @overload # type: ignore[override] + @overload def __radd__(self: Index[Never], other: _str) -> Never: ... @overload def __radd__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index ef3c625ff..be005260b 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1876,7 +1876,7 @@ class Series(IndexOpsMixin[S1], NDFrame): fill_value: float | None = None, axis: int = 0, ) -> Series[_str]: ... - @overload # type: ignore[override] + @overload def __radd__(self: Series[Never], other: _str) -> Never: ... @overload def __radd__( diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 280b57261..116155f32 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -262,12 +262,10 @@ def test_types_to_numpy() -> None: def test_index_arithmetic() -> None: # GH 287 idx = pd.Index([1, 2.2, 3], dtype=float) - check(assert_type(idx + 3, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(idx - 3, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(idx * 3, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(idx / 3, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(idx // 3, "pd.Index[float]"), pd.Index, np.float64) - check(assert_type(3 + idx, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(3 - idx, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(3 * idx, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(3 / idx, "pd.Index[float]"), pd.Index, np.float64)