ENH: Forward operators to underlying types #820
Description
Describe the bug
Series is a container that forwards operators (+, -, ...) to the underlying types.
Some of the operators of Series could be better, for example, __add__ which:
- returns
Series[Unknown]instead ofSeries[return type] - or uses non-exiting helper classes to restrict the types
Instead of returning too wide types, enumerating all typing combinations, or creating non-existing helper classes, it might(?) be possible to abuse Protocols to forward operators to the underlying types.
To Reproduce
The following code works with both mypy and pyright
from __future__ import annotations
from typing import Protocol, TypeVar, overload, Generic
_T = TypeVar("_T", covariant=True)
_Other = TypeVar("_Other", contravariant=True)
_Result = TypeVar("_Result", covariant=True)
class add(Protocol[_Other, _Result]):
def __add__(self, other: _Other, /) -> _Result:
...
class radd(Protocol[_Other, _Result]):
def __radd__(self, other: _Other, /) -> _Result:
...
class Series(Generic[_T]): # simulate Series with one element
def __init__(self, x: _T) -> None:
self.x = x
# Series + Series
@overload
def __add__(
self: Series[add[_Other, _Result]], other: Series[_Other], /
) -> Series[_Result]:
...
@overload
def __add__(self, other: Series[radd[_T, _Result]], /) -> Series[_Result]:
...
# Series + scalar
@overload
def __add__(
self: Series[add[_Other, _Result]], other: _Other, /
) -> Series[_Result]:
...
@overload
def __add__(self, other: radd[_T, _Result], /) -> Series[_Result]:
...
def __add__(self, other):
if isinstance(other, Series):
return Series(self.x + other.x)
return Series(self.x + other)
def test_int(x: Series[int], y: int):
reveal_type(x + y) # Series[int] (pandas-stubs: Series[int])
def test_float(x: Series[int], y: float):
reveal_type(x + y) # Series[float] (pandas-stubs: Series[Unknown])
def test_series_int(x: Series[int], y: Series[int]):
reveal_type(x + y) # Series[int] (pandas-stubs: Series[int])
def test_series_float(x: Series[int], y: Series[float]):
reveal_type(x + y) # Series[float] (pandas-stubs: Series[Unknown])
def test_error(x: Series[int], y: str):
x + y # error :) (pandas-stubs: no error)
def test_series_error(x: Series[int], y: Series[str]):
x + y # error :) (pandas-stubs: no error)