diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 9cdbbf91..db9d4919 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -25,6 +25,7 @@ from typing import ( Generic, Literal, NoReturn, + TypeVar, overload, ) @@ -174,6 +175,8 @@ from pandas._typing import ( VoidDtypeArg, WriteBuffer, np_ndarray_anyint, + np_ndarray_complex, + np_ndarray_float, npt, num, ) @@ -183,6 +186,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 @@ -1580,6 +1585,51 @@ class Series(IndexOpsMixin[S1], NDFrame): # just failed to generate these so I couldn't match # them up. @overload + def __add__( + self: Series[Never], + other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64, + ) -> Series: ... + @overload + def __add__(self: Series[int], 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: Series[Never]) -> Series: ... + @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: Series[Never]) -> Series: ... + @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: ... @overload def __add__( @@ -1626,6 +1676,24 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ... @overload + def __radd__( + self: Series[Never], other: num | _str | _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: _T_COMPLEX | Sequence[_T_COMPLEX] + ) -> Series[complex]: ... + @overload def __radd__(self, other: S1 | Series[S1]) -> Self: ... @overload def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ... @@ -1703,6 +1771,92 @@ class Series(IndexOpsMixin[S1], NDFrame): @property def loc(self) -> _LocIndexerSeries[S1]: ... # Methods + @overload + def add( + self: Series[Never], + other: ( + Sequence + | np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + | Series + | Scalar + ), + 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, @@ -1945,6 +2099,92 @@ class Series(IndexOpsMixin[S1], NDFrame): min_count: int = ..., **kwargs: Any, ) -> Scalar: ... + @overload + def radd( + self: Series[Never], + other: ( + Sequence + | np_ndarray_anyint + | np_ndarray_float + | np_ndarray_complex + | Series + | Scalar + ), + 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, diff --git a/pyproject.toml b/pyproject.toml index 29239c33..411a8364 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,11 +35,11 @@ types-pytz = ">= 2022.1.1" numpy = ">= 1.23.5" [tool.poetry.group.dev.dependencies] -mypy = "1.16.0" +mypy = ">=1.17.0" pandas = "2.3.0" 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/__init__.py b/tests/__init__.py index e0682a2d..c4a94e77 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -64,7 +64,7 @@ def check( if not isinstance(actual, klass): raise RuntimeError(f"Expected type '{klass}' but got '{type(actual)}'") if dtype is None: - return actual # type: ignore[return-value] + return actual if isinstance(actual, pd.Series): value = actual.iloc[index_to_check_for_type] @@ -82,7 +82,7 @@ def check( if not isinstance(value, dtype): raise RuntimeError(f"Expected type '{dtype}' but got '{type(value)}'") - return actual # type: ignore[return-value] + return actual def pytest_warns_bounded( 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..3b0f910b --- /dev/null +++ b/tests/series/arithmetic/complex/test_add.py @@ -0,0 +1,95 @@ +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) + + # check(assert_type(i + l, "pd.Series[complex]"), pd.Series, np.complex128) + # check(assert_type(f + l, "pd.Series[complex]"), pd.Series, np.complex128) + # check(assert_type(c + l, "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_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..4ecdd744 --- /dev/null +++ b/tests/series/arithmetic/float/test_add.py @@ -0,0 +1,95 @@ +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) + + # check(assert_type(i + l, "pd.Series[float]"), pd.Series, np.float64) + # check(assert_type(f + l, "pd.Series[float]"), pd.Series, np.float64) + # check(assert_type(c + l, "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_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..bb4aec79 --- /dev/null +++ b/tests/series/arithmetic/int/test_add.py @@ -0,0 +1,95 @@ +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) + + # check(assert_type(i + l, "pd.Series[int]"), pd.Series, np.int64) + # check(assert_type(f + l, "pd.Series[float]"), pd.Series, np.float64) + # check(assert_type(c + l, "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_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..e1d054d2 --- /dev/null +++ b/tests/series/arithmetic/test_add.py @@ -0,0 +1,95 @@ +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) + + # check(assert_type(i + l, pd.Series), pd.Series) + # check(assert_type(f + l, pd.Series), pd.Series) + # check(assert_type(c + l, 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_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