diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 3a3ba8c4..631709b8 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,23 @@ 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) -> Series: ... + @overload + def __radd__( + self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] + ) -> Series[_T_COMPLEX]: ... + @overload + def __radd__(self: Series[float], other: int | Sequence[int]) -> Series[float]: ... + @overload + def __radd__( + self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX] + ) -> Series[_T_COMPLEX]: ... + @overload + def __radd__( + self: Series[complex], other: complex | Sequence[complex] + ) -> Series[complex]: ... @overload - def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ... + def __radd__(self, other: S1 | Series[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 +1796,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 +2119,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, 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..5c43763e --- /dev/null +++ b/tests/series/arithmetic/complex/test_add.py @@ -0,0 +1,114 @@ +import numpy as np +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 ndarray instead of `pd.Series[...]` in reality, which we cannot fix + check( + assert_type( # type: ignore[assert-type] + i + left, "pd.Series[complex]" # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + np.complex128, + ) + check( + assert_type( # type: ignore[assert-type] + f + left, "pd.Series[complex]" # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + np.complex128, + ) + check( + assert_type( # type: ignore[assert-type] + c + left, "pd.Series[complex]" # pyright: ignore[reportAssertTypeFailure] + ), + 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..16085f36 --- /dev/null +++ b/tests/series/arithmetic/float/test_add.py @@ -0,0 +1,114 @@ +import numpy as np +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 ndarray instead of `pd.Series[...]` in reality, which we cannot fix + check( + assert_type( # type: ignore[assert-type] + i + left, "pd.Series[float]" # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + np.float64, + ) + check( + assert_type( # type: ignore[assert-type] + f + left, "pd.Series[float]" # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + np.float64, + ) + check( + assert_type( # type: ignore[assert-type] + c + left, "pd.Series[complex]" # pyright: ignore[reportAssertTypeFailure] + ), + 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..cb78181c --- /dev/null +++ b/tests/series/arithmetic/int/test_add.py @@ -0,0 +1,114 @@ +import numpy as np +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 ndarray instead of `pd.Series[...]` in reality, which we cannot fix + check( + assert_type( # type: ignore[assert-type] + i + left, "pd.Series[int]" # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + np.int64, + ) + check( + assert_type( # type: ignore[assert-type] + f + left, "pd.Series[float]" # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + np.float64, + ) + check( + assert_type( # type: ignore[assert-type] + c + left, "pd.Series[complex]" # pyright: ignore[reportAssertTypeFailure] + ), + 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..0001f6ec --- /dev/null +++ b/tests/series/arithmetic/test_add.py @@ -0,0 +1,111 @@ +import numpy as np +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 ndarray instead of `pd.Series[...]` in reality, which we cannot fix + check( + assert_type( # type: ignore[assert-type] + i + left, pd.Series # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + ) + check( + assert_type( # type: ignore[assert-type] + f + left, pd.Series # pyright: ignore[reportAssertTypeFailure] + ), + pd.Series, + ) + check( + assert_type( # type: ignore[assert-type] + c + left, pd.Series # pyright: ignore[reportAssertTypeFailure] + ), + 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_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