diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 3a3ba8c4..f9afa811 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -25,6 +25,7 @@ from typing import ( Generic, Literal, NoReturn, + TypeVar, final, overload, ) @@ -175,6 +176,8 @@ from pandas._typing import ( VoidDtypeArg, WriteBuffer, np_ndarray_anyint, + np_ndarray_complex, + np_ndarray_float, npt, num, ) @@ -184,6 +187,8 @@ from pandas.core.dtypes.dtypes import CategoricalDtype from pandas.plotting import PlotAccessor +_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex) + class _iLocIndexerSeries(_iLocIndexer, Generic[S1]): # get item @overload @@ -1617,12 +1622,48 @@ class Series(IndexOpsMixin[S1], NDFrame): # just failed to generate these so I couldn't match # them up. @overload - def __add__(self, other: S1 | Self) -> Self: ... - @overload def __add__( - self, - other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64, + self: Series[Never], + other: Scalar | _ListLike | Series, ) -> Series: ... + @overload + def __add__(self, other: Series[Never]) -> Series: ... + @overload + def __add__( + self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX] + ) -> Series[_T_COMPLEX]: ... + @overload + def __add__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ... + @overload + def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... + @overload + def __add__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ... + @overload + def __add__( + self: Series[float], + other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int], + ) -> Series[float]: ... + @overload + def __add__( + self: Series[float], + 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] + | Series[_T_COMPLEX] + | np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + ), + ) -> Series[complex]: ... + @overload + def __add__(self, other: S1 | Self) -> Self: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] def __and__( # pyright: ignore[reportOverlappingOverload] @@ -1663,9 +1704,40 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ... @overload - def __radd__(self, other: S1 | Series[S1]) -> Self: ... + def __radd__(self: Series[Never], other: Scalar | _ListLike) -> Series: ... + @overload + def __radd__( + self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] + ) -> Series[_T_COMPLEX]: ... + @overload + def __radd__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ... + @overload + def __radd__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... + @overload + def __radd__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ... + @overload + def __radd__( + self: Series[float], + other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float, + ) -> Series[float]: ... + @overload + def __radd__( + self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX] + ) -> Series[_T_COMPLEX]: ... @overload - def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ... + def __radd__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ... + @overload + def __radd__( + self: Series[complex], + other: ( + np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + | Sequence[_T_COMPLEX] + ), + ) -> Series[complex]: ... + @overload + def __radd__(self, other: S1) -> Self: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] def __rand__( # pyright: ignore[reportOverlappingOverload] @@ -1741,13 +1813,92 @@ class Series(IndexOpsMixin[S1], NDFrame): @property def loc(self) -> _LocIndexerSeries[S1]: ... # Methods + @overload + def add( + self: Series[Never], + other: Scalar | _ListLike | Series, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series: ... + @overload + def add( + self: Series[int], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[_T_COMPLEX]: ... + @overload + def add( + self: Series[int], + other: np_ndarray_anyint, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[int]: ... + @overload + def add( + self: Series[int], + other: np_ndarray_float, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[float]: ... + @overload + def add( + self: Series[int], + other: np_ndarray_complex, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload + def add( + self: Series[float], + other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[float]: ... + @overload + def add( + self: Series[float], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[_T_COMPLEX]: ... + @overload + def add( + self: Series[float], + other: np_ndarray_complex, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload + def add( + self: Series[complex], + other: ( + Sequence[_T_COMPLEX] + | np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + | Series[_T_COMPLEX] + ), + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload def add( self, - other: Series[S1] | Scalar, + other: S1 | Series[S1], level: Level | None = ..., fill_value: float | None = ..., axis: int = ..., - ) -> Series[S1]: ... + ) -> Self: ... def all( self, axis: AxisIndex = ..., @@ -1985,13 +2136,92 @@ class Series(IndexOpsMixin[S1], NDFrame): min_count: int = ..., **kwargs: Any, ) -> Scalar: ... + @overload + def radd( + self: Series[Never], + other: Scalar | _ListLike | Series, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex = ..., + ) -> Series: ... + @overload + def radd( + self: Series[int], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[_T_COMPLEX]: ... + @overload + def radd( + self: Series[int], + other: np_ndarray_anyint, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[int]: ... + @overload + def radd( + self: Series[int], + other: np_ndarray_float, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[float]: ... + @overload + def radd( + self: Series[int], + other: np_ndarray_complex, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload + def radd( + self: Series[float], + other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[float]: ... + @overload + def radd( + self: Series[float], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[_T_COMPLEX]: ... + @overload + def radd( + self: Series[float], + other: np_ndarray_complex, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload + def radd( + self: Series[complex], + other: ( + Sequence[_T_COMPLEX] + | np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + | Series[_T_COMPLEX] + ), + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload def radd( self, - other: Series[S1] | Scalar, + other: S1 | Series[S1], level: Level | None = ..., fill_value: float | None = ..., axis: AxisIndex = ..., - ) -> Series[S1]: ... + ) -> Self: ... def rdivmod( self, other: Series[S1] | Scalar, @@ -2395,7 +2625,7 @@ class PeriodSeries(Series[Period]): ) -> Never: ... class OffsetSeries(Series[BaseOffset]): - @overload # type: ignore[override] + @overload def __radd__(self, other: Period) -> PeriodSeries: ... @overload def __radd__( # pyright: ignore[reportIncompatibleMethodOverride] diff --git a/pyproject.toml b/pyproject.toml index a60a4f6b..0de7fe6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ mypy = "1.17.0" pandas = "2.3.1" pyarrow = ">=10.0.1" pytest = ">=7.1.2" -pyright = ">=1.1.400" +pyright = ">=1.1.403" ty = "^0.0.1a8" pyrefly = "^0.21.0" poethepoet = ">=0.16.5" diff --git a/tests/series/__init__.py b/tests/series/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/series/arithmetic/__init__.py b/tests/series/arithmetic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/series/arithmetic/complex/__init__.py b/tests/series/arithmetic/complex/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/series/arithmetic/complex/test_add.py b/tests/series/arithmetic/complex/test_add.py new file mode 100644 index 00000000..2ab7d1af --- /dev/null +++ b/tests/series/arithmetic/complex/test_add.py @@ -0,0 +1,99 @@ +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 = pd.Series([1j, 2j, 3j]) # left operand + + +def test_add_py_scalar() -> None: + """Test pd.Series[complex] + Python native scalars""" + i, f, c = 1, 1.0, 1j + + check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_py_sequence() -> None: + """Test pd.Series[complex] + Python native sequence""" + i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_numpy_array() -> None: + """Test pd.Series[complex] + numpy array""" + 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, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Series`s with the correct element type. + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.complex128) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.complex128) + check(assert_type(c + left, "npt.NDArray[np.complex128]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_pd_series() -> None: + """Test pd.Series[complex] + pandas series""" + i = pd.Series([2, 3, 5]) + f = pd.Series([1.0, 2.0, 3.0]) + c = pd.Series([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) diff --git a/tests/series/arithmetic/float/__init__.py b/tests/series/arithmetic/float/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/series/arithmetic/float/test_add.py b/tests/series/arithmetic/float/test_add.py new file mode 100644 index 00000000..61b77742 --- /dev/null +++ b/tests/series/arithmetic/float/test_add.py @@ -0,0 +1,99 @@ +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 = pd.Series([1.0, 2.0, 3.0]) # left operand + + +def test_add_py_scalar() -> None: + """Test pd.Series[float] + Python native scalars""" + i, f, c = 1, 1.0, 1j + + check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_py_sequence() -> None: + """Test pd.Series[float] + Python native sequence""" + i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_numpy_array() -> None: + """Test pd.Series[float] + numpy array""" + 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, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Series`s with the correct element type. + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.float64) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.float64) + check(assert_type(c + left, "npt.NDArray[np.complex128]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_pd_series() -> None: + """Test pd.Series[float] + pandas series""" + i = pd.Series([2, 3, 5]) + f = pd.Series([1.0, 2.0, 3.0]) + c = pd.Series([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) diff --git a/tests/series/arithmetic/int/__init__.py b/tests/series/arithmetic/int/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/series/arithmetic/int/test_add.py b/tests/series/arithmetic/int/test_add.py new file mode 100644 index 00000000..ad25aeb3 --- /dev/null +++ b/tests/series/arithmetic/int/test_add.py @@ -0,0 +1,103 @@ +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 = pd.Series([1, 2, 3]) # left operand + + +def test_add_py_scalar() -> None: + """Test pd.Series[int] + Python native scalars""" + i, f, c = 1, 1.0, 1j + + check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_py_sequence() -> None: + """Test pd.Series[int] + Python native sequence""" + i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_numpy_array() -> None: + """Test pd.Series[int] + numpy array""" + 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, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + # `numpy` typing gives the corresponding `ndarray`s in the static type + # checking, where our `__radd__` cannot override. At runtime, they return + # `Series`s with the correct element type. + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.int64) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.float64) + check( + assert_type(c + left, "npt.NDArray[np.complex128]"), + pd.Series, + np.complex128, + ) + + check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) + + +def test_add_pd_series() -> None: + """Test pd.Series[int] + pandas series""" + i = pd.Series([2, 3, 5]) + f = pd.Series([1.0, 2.0, 3.0]) + c = pd.Series([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128) + + check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128) diff --git a/tests/series/arithmetic/test_add.py b/tests/series/arithmetic/test_add.py new file mode 100644 index 00000000..d5407dc7 --- /dev/null +++ b/tests/series/arithmetic/test_add.py @@ -0,0 +1,107 @@ +import numpy as np +from numpy import typing as npt # noqa: F401 +from numpy._typing import _32Bit # noqa: F401 +import pandas as pd +from typing_extensions import assert_type + +from tests import check + +left = pd.DataFrame({"a": [1, 2, 3]})["a"] # left operand + + +def test_add_py_scalar() -> None: + """Test pd.Series[Any] + Python native scalars""" + i, f, c = 1, 1.0, 1j + + check(assert_type(left + i, pd.Series), pd.Series) + check(assert_type(left + f, pd.Series), pd.Series) + check(assert_type(left + c, pd.Series), pd.Series) + + check(assert_type(i + left, pd.Series), pd.Series) + check(assert_type(f + left, pd.Series), pd.Series) + check(assert_type(c + left, pd.Series), pd.Series) + + check(assert_type(left.add(i), pd.Series), pd.Series) + check(assert_type(left.add(f), pd.Series), pd.Series) + check(assert_type(left.add(c), pd.Series), pd.Series) + + check(assert_type(left.radd(i), pd.Series), pd.Series) + check(assert_type(left.radd(f), pd.Series), pd.Series) + check(assert_type(left.radd(c), pd.Series), pd.Series) + + +def test_add_py_sequence() -> None: + """Test pd.Series[Any] + Python native sequence""" + i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j] + + check(assert_type(left + i, pd.Series), pd.Series) + check(assert_type(left + f, pd.Series), pd.Series) + check(assert_type(left + c, pd.Series), pd.Series) + + check(assert_type(i + left, pd.Series), pd.Series) + check(assert_type(f + left, pd.Series), pd.Series) + check(assert_type(c + left, pd.Series), pd.Series) + + check(assert_type(left.add(i), pd.Series), pd.Series) + check(assert_type(left.add(f), pd.Series), pd.Series) + check(assert_type(left.add(c), pd.Series), pd.Series) + + check(assert_type(left.radd(i), pd.Series), pd.Series) + check(assert_type(left.radd(f), pd.Series), pd.Series) + check(assert_type(left.radd(c), pd.Series), pd.Series) + + +def test_add_numpy_array() -> None: + """Test pd.Series[Any] + numpy array""" + 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.complex64) + + check(assert_type(left + i, pd.Series), pd.Series) + check(assert_type(left + f, pd.Series), pd.Series) + check(assert_type(left + c, pd.Series), pd.Series) + + # `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(i + left, "npt.NDArray[np.int64]"), pd.Series # type: ignore[assert-type] + ) + check( + assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series # type: ignore[assert-type] + ) + check( + assert_type(c + left, "npt.NDArray[np.complexfloating[_32Bit, _32Bit]]"), pd.Series # type: ignore[assert-type] + ) + + check(assert_type(left.add(i), pd.Series), pd.Series) + check(assert_type(left.add(f), pd.Series), pd.Series) + check(assert_type(left.add(c), pd.Series), pd.Series) + + check(assert_type(left.radd(i), pd.Series), pd.Series) + check(assert_type(left.radd(f), pd.Series), pd.Series) + check(assert_type(left.radd(c), pd.Series), pd.Series) + + +def test_add_pd_series() -> None: + """Test pd.Series[Any] + pandas series""" + i = pd.Series([2, 3, 5]) + f = pd.Series([1.0, 2.0, 3.0]) + c = pd.Series([1.1j, 2.2j, 4.1j]) + + check(assert_type(left + i, pd.Series), pd.Series) + check(assert_type(left + f, pd.Series), pd.Series) + check(assert_type(left + c, pd.Series), pd.Series) + + check(assert_type(i + left, pd.Series), pd.Series) + check(assert_type(f + left, pd.Series), pd.Series) + check(assert_type(c + left, pd.Series), pd.Series) + + check(assert_type(left.add(i), pd.Series), pd.Series) + check(assert_type(left.add(f), pd.Series), pd.Series) + check(assert_type(left.add(c), pd.Series), pd.Series) + + check(assert_type(left.radd(i), pd.Series), pd.Series) + check(assert_type(left.radd(f), pd.Series), pd.Series) + check(assert_type(left.radd(c), pd.Series), pd.Series) diff --git a/tests/test_series.py b/tests/series/test_series.py similarity index 100% rename from tests/test_series.py rename to tests/series/test_series.py