ENH: Static typing support for custom array containers #21737
Description
Proposed new feature or change:
It would be nice to have static typing support for custom array containers like pandas Series, xarrays DataArray (pydata/xarray#6524) etc.
Such that you can do:
import pandas as pd
import numpy as np
x = pd.Series([1, 2, 3])
y = np.exp(x)
reveal_type(y) # Should be pd.Series but is np.ndarrayI assume this should be possible using Protocols and using something like this:
from typing import Protocol, TypeVar
class HasArrayUFunc(Protocol):
def __array_ufunc__(ufunc, method, *inputs, **kwargs):
pass
ArrayOrHasArrayUFunc = TypeVar("ArrayOrHasArrayUFunc", ndarray, HasArrayUFunc)
def exp(x: ArrayOrHasArrayUFunc) -> ArrayOrHasArrayUFunc: ...