From e6b97c1cc84061aba734763f2a4d7a232d80e6a6 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Fri, 8 Aug 2025 10:29:36 +0200 Subject: [PATCH] feat: addition for bools --- pandas-stubs/core/series.pyi | 182 ++++++++++++++++++-- tests/series/arithmetic/bool/__init__.py | 0 tests/series/arithmetic/bool/test_add.py | 129 ++++++++++++++ tests/series/arithmetic/complex/test_add.py | 30 +++- tests/series/arithmetic/float/test_add.py | 22 ++- tests/series/arithmetic/int/test_add.py | 22 ++- tests/series/arithmetic/test_add.py | 29 +++- 7 files changed, 387 insertions(+), 27 deletions(-) create mode 100644 tests/series/arithmetic/bool/__init__.py create mode 100644 tests/series/arithmetic/bool/test_add.py diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index e4006c266..d206adcb6 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -176,6 +176,7 @@ from pandas._typing import ( VoidDtypeArg, WriteBuffer, np_ndarray_anyint, + np_ndarray_bool, np_ndarray_complex, np_ndarray_float, npt, @@ -187,6 +188,7 @@ 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]): @@ -1619,10 +1621,29 @@ class Series(IndexOpsMixin[S1], NDFrame): def __add__(self, other: Series[Never]) -> Series: ... @overload def __add__( - self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX] + self: Series[bool], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], ) -> Series[_T_COMPLEX]: ... @overload - def __add__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ... + def __add__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... + @overload + def __add__(self: Series[bool], other: np_ndarray_anyint) -> Series[int]: ... + @overload + def __add__(self: Series[bool], other: np_ndarray_float) -> Series[float]: ... + @overload + def __add__(self: Series[bool], other: np_ndarray_complex) -> Series[complex]: ... + @overload + def __add__( + self: Series[int], + other: ( + bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + ), + ) -> Series[int]: ... + @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_float) -> Series[float]: ... @overload @@ -1630,7 +1651,14 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __add__( self: Series[float], - other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int], + other: ( + _T_INT + | Sequence[_T_INT] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | Series[_T_INT] + ), ) -> Series[float]: ... @overload def __add__( @@ -1643,12 +1671,12 @@ class Series(IndexOpsMixin[S1], NDFrame): def __add__( self: Series[complex], other: ( - _T_COMPLEX - | Sequence[_T_COMPLEX] - | Series[_T_COMPLEX] + Sequence[_T_COMPLEX] + | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex + | Series[_T_COMPLEX] ), ) -> Series[complex]: ... @overload @@ -1663,7 +1691,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series: ... @overload def add( - self: Series[int], + self: Series[bool], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], level: Level | None = None, fill_value: float | None = None, @@ -1671,13 +1699,55 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[_T_COMPLEX]: ... @overload def add( - self: Series[int], + self: Series[bool], + other: np_ndarray_bool, + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[bool]: ... + @overload + def add( + self: Series[bool], other: np_ndarray_anyint, level: Level | None = None, fill_value: float | None = None, axis: int = 0, ) -> Series[int]: ... @overload + def add( + self: Series[bool], + other: np_ndarray_float, + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[float]: ... + @overload + def add( + self: Series[bool], + other: np_ndarray_complex, + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[complex]: ... + @overload + def add( + self: Series[int], + other: ( + bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + ), + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[int]: ... + @overload + def add( + self: Series[int], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[_T_COMPLEX]: ... + @overload def add( self: Series[int], other: np_ndarray_float, @@ -1696,7 +1766,14 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def add( self: Series[float], - other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int], + other: ( + _T_INT + | Sequence[_T_INT] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | Series[_T_INT] + ), level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1722,6 +1799,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self: Series[complex], other: ( Sequence[_T_COMPLEX] + | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex @@ -1743,10 +1821,25 @@ class Series(IndexOpsMixin[S1], NDFrame): def __radd__(self: Series[Never], other: Scalar | _ListLike) -> Series: ... @overload def __radd__( - self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] + self: Series[bool], other: _T_COMPLEX | Sequence[_T_COMPLEX] ) -> Series[_T_COMPLEX]: ... @overload - def __radd__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ... + def __radd__(self: Series[bool], other: np_ndarray_bool) -> Series[bool]: ... + @overload + def __radd__(self: Series[bool], other: np_ndarray_anyint) -> Series[int]: ... + @overload + def __radd__(self: Series[bool], other: np_ndarray_float) -> Series[float]: ... + @overload + def __radd__(self: Series[bool], other: np_ndarray_complex) -> Series[complex]: ... + @overload + def __radd__( + self: Series[int], + other: bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint, + ) -> Series[int]: ... + @overload + def __radd__( + self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] + ) -> Series[_T_COMPLEX]: ... @overload def __radd__(self: Series[int], other: np_ndarray_float) -> Series[float]: ... @overload @@ -1754,7 +1847,13 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __radd__( self: Series[float], - other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float, + other: ( + _T_INT + | Sequence[_T_INT] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + ), ) -> Series[float]: ... @overload def __radd__( @@ -1766,7 +1865,8 @@ class Series(IndexOpsMixin[S1], NDFrame): def __radd__( self: Series[complex], other: ( - np_ndarray_anyint + np_ndarray_bool + | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex | Sequence[_T_COMPLEX] @@ -1784,7 +1884,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series: ... @overload def radd( - self: Series[int], + self: Series[bool], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], level: Level | None = None, fill_value: float | None = None, @@ -1792,13 +1892,55 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[_T_COMPLEX]: ... @overload def radd( - self: Series[int], + self: Series[bool], + other: np_ndarray_bool, + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[bool]: ... + @overload + def radd( + self: Series[bool], other: np_ndarray_anyint, level: Level | None = None, fill_value: float | None = None, axis: int = 0, ) -> Series[int]: ... @overload + def radd( + self: Series[bool], + other: np_ndarray_float, + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[float]: ... + @overload + def radd( + self: Series[bool], + other: np_ndarray_complex, + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[complex]: ... + @overload + def radd( + self: Series[int], + other: ( + bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Series[bool] + ), + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[int]: ... + @overload + def radd( + self: Series[int], + other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX], + level: Level | None = None, + fill_value: float | None = None, + axis: int = 0, + ) -> Series[_T_COMPLEX]: ... + @overload def radd( self: Series[int], other: np_ndarray_float, @@ -1817,7 +1959,14 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def radd( self: Series[float], - other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int], + other: ( + _T_INT + | Sequence[_T_INT] + | np_ndarray_bool + | np_ndarray_anyint + | np_ndarray_float + | Series[_T_INT] + ), level: Level | None = None, fill_value: float | None = None, axis: int = 0, @@ -1843,6 +1992,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self: Series[complex], other: ( Sequence[_T_COMPLEX] + | np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex diff --git a/tests/series/arithmetic/bool/__init__.py b/tests/series/arithmetic/bool/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/series/arithmetic/bool/test_add.py b/tests/series/arithmetic/bool/test_add.py new file mode 100644 index 000000000..d2a77d55f --- /dev/null +++ b/tests/series/arithmetic/bool/test_add.py @@ -0,0 +1,129 @@ +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([True, True, False]) # left operand + + +def test_add_py_scalar() -> None: + """Test pd.Series[bool] + Python native scalars""" + b, i, f, c = True, 1, 1.0, 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 + ) + + +def test_add_py_sequence() -> None: + """Test pd.Series[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.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 + ) + + +def test_add_numpy_array() -> None: + """Test pd.Series[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.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) + + # `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(b + left, "npt.NDArray[np.bool_]"), pd.Series, np.bool_) + check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.integer) + check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.floating) + check( + assert_type(c + left, "npt.NDArray[np.complex128]"), + 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 + ) + + +def test_add_pd_series() -> None: + """Test pd.Series[bool] + pandas series""" + b = pd.Series([True, False, True]) + 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 + 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 269204087..e3668b338 100644 --- a/tests/series/arithmetic/complex/test_add.py +++ b/tests/series/arithmetic/complex/test_add.py @@ -10,20 +10,26 @@ def test_add_py_scalar() -> None: """Test pd.Series[complex] + Python native scalars""" - i, f, c = 1, 1.0, 1j + b, i, f, c = True, 1, 1.0, 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 ) @@ -37,20 +43,26 @@ def test_add_py_scalar() -> None: 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] + 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.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 ) @@ -64,10 +76,12 @@ def test_add_py_sequence() -> None: def test_add_numpy_array() -> None: """Test pd.Series[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.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) @@ -75,6 +89,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 with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Series, np.complexfloating) check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.complexfloating) check( assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.complexfloating @@ -85,10 +100,14 @@ def test_add_numpy_array() -> None: 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 ) @@ -102,22 +121,29 @@ def test_add_numpy_array() -> None: def test_add_pd_series() -> None: """Test pd.Series[complex] + pandas series""" + b = pd.Series([True, False, True]) 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 + 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 ) diff --git a/tests/series/arithmetic/float/test_add.py b/tests/series/arithmetic/float/test_add.py index 6864c596e..41fc6b96b 100644 --- a/tests/series/arithmetic/float/test_add.py +++ b/tests/series/arithmetic/float/test_add.py @@ -10,20 +10,24 @@ def test_add_py_scalar() -> None: """Test pd.Series[float] + Python native scalars""" - i, f, c = 1, 1.0, 1j + b, i, f, c = True, 1, 1.0, 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( @@ -33,20 +37,24 @@ def test_add_py_scalar() -> None: 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] + 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.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( @@ -56,10 +64,12 @@ def test_add_py_sequence() -> None: def test_add_numpy_array() -> None: """Test pd.Series[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.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) @@ -67,6 +77,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 with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Series, np.floating) check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.floating) check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.floating) check( @@ -75,10 +86,12 @@ def test_add_numpy_array() -> None: 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( @@ -88,22 +101,27 @@ def test_add_numpy_array() -> None: def test_add_pd_series() -> None: """Test pd.Series[float] + pandas series""" + b = pd.Series([True, False, True]) 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 + 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( diff --git a/tests/series/arithmetic/int/test_add.py b/tests/series/arithmetic/int/test_add.py index 2117dcd90..b68b5e235 100644 --- a/tests/series/arithmetic/int/test_add.py +++ b/tests/series/arithmetic/int/test_add.py @@ -10,20 +10,24 @@ def test_add_py_scalar() -> None: """Test pd.Series[int] + Python native scalars""" - i, f, c = 1, 1.0, 1j + b, i, f, c = True, 1, 1.0, 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( @@ -33,20 +37,24 @@ def test_add_py_scalar() -> None: 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] + 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.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( @@ -56,10 +64,12 @@ def test_add_py_sequence() -> None: def test_add_numpy_array() -> None: """Test pd.Series[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.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) @@ -67,6 +77,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 with the correct element type. + check(assert_type(b + left, "npt.NDArray[np.bool_]"), pd.Series, np.integer) check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.integer) check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.floating) check( @@ -75,10 +86,12 @@ def test_add_numpy_array() -> None: 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( @@ -88,22 +101,27 @@ def test_add_numpy_array() -> None: def test_add_pd_series() -> None: """Test pd.Series[int] + pandas series""" + b = pd.Series([True, False, True]) 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 + 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( diff --git a/tests/series/arithmetic/test_add.py b/tests/series/arithmetic/test_add.py index d5407dc76..e4d68c94f 100644 --- a/tests/series/arithmetic/test_add.py +++ b/tests/series/arithmetic/test_add.py @@ -1,6 +1,5 @@ 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 @@ -11,20 +10,24 @@ def test_add_py_scalar() -> None: """Test pd.Series[Any] + Python native scalars""" - i, f, c = 1, 1.0, 1j + b, i, f, c = True, 1, 1.0, 1j + check(assert_type(left + b, pd.Series), pd.Series) 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(b + left, 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(b), 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(b), 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) @@ -32,20 +35,24 @@ def test_add_py_scalar() -> None: 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] + 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.Series), pd.Series) 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(b + left, 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(b), 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(b), 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) @@ -53,10 +60,12 @@ def test_add_py_sequence() -> None: def test_add_numpy_array() -> None: """Test pd.Series[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.complex64) + c = np.array([1.1j, 2.2j, 4.1j], np.complex128) + check(assert_type(left + b, pd.Series), pd.Series) 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) @@ -65,6 +74,9 @@ def test_add_numpy_array() -> None: # 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, "npt.NDArray[np.bool_]"), pd.Series # type: ignore[assert-type] + ) check( assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series # type: ignore[assert-type] ) @@ -72,13 +84,15 @@ def test_add_numpy_array() -> None: 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] + assert_type(c + left, "npt.NDArray[np.complex128]"), pd.Series # type: ignore[assert-type] ) + check(assert_type(left.add(b), 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(b), 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) @@ -86,22 +100,27 @@ def test_add_numpy_array() -> None: def test_add_pd_series() -> None: """Test pd.Series[Any] + pandas series""" + b = pd.Series([True, False, True]) 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 + b, pd.Series), pd.Series) 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(b + left, 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(b), 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(b), 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)