Skip to content
767 changes: 491 additions & 276 deletions pandas-stubs/core/series.pyi

Large diffs are not rendered by default.

136 changes: 83 additions & 53 deletions tests/series/arithmetic/complex/test_add.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,42 +12,54 @@ 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(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(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(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(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.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(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)
check(
assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating
)


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(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(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(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(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.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(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)
check(
assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating
)


def test_add_numpy_array() -> None:
Expand All @@ -56,24 +68,36 @@ def test_add_numpy_array() -> None:
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(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)

# `numpy` typing gives the corresponding `ndarray`s in the static type
# checking, where our `__radd__` cannot override. At runtime, they return
# `Series`s with the correct element type.
check(assert_type(i + left, "npt.NDArray[np.int64]"), pd.Series, np.complex128)
check(assert_type(f + left, "npt.NDArray[np.float64]"), pd.Series, np.complex128)
check(assert_type(c + left, "npt.NDArray[np.complex128]"), pd.Series, np.complex128)

check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128)
check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128)
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)

check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128)
check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128)
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
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
)
check(
assert_type(c + left, "npt.NDArray[np.complex128]"),
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(i), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating
)


def test_add_pd_series() -> None:
Expand All @@ -82,18 +106,24 @@ def test_add_pd_series() -> None:
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)
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(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(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(i), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complexfloating
)
check(
assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complexfloating
)
Loading
Loading