If I want to add type annotation to the return value of the following function, how should I?

import numpy as np

def myfunc(x: np.ndarray):
    return x.sum()

The type of return value of the function vary:

• np.int8 if x.dtype is dtype('int8')
• np.int16 if x.dtype is dtype('int16')
• np.int32 if x.dtype is dtype('int32')
• np.float16 if x.dtype is dtype('float16')
• np.float32 if x.dtype is dtype('float32') ... etc.

Of cause I want to avoid using -> Any.

CodePudding user response：

You can use NDArray from the numpy typing module and define a TypeVar that is bound to numpy scalar types. This bound type you can pass to the NDArray type and annotate the return type of your function with that same type.

import numpy as np
from typing import TypeVar
from numpy.typing import NDArray

ScalarType = TypeVar("ScalarType", bound=np.generic, covariant=True)

def myfunc(x: NDArray[ScalarType]) -> ScalarType:
    return x.sum()