I wonder what’s the right way to type annotate a python function as below?
def f(value): if isinstance(value, str): return None return value
This is a toy example that is analogous to a python I encountered that needs type annotated. Using Union
like
def f(value: Union[float, int, str]) -> Union[float, int, None]: if isinstance(value, str): return None return value
doesn’t feel right because it doesn’t enforce the rules
- an
int
input must result in anint
output. - a
float
input must result in afloat
output. - a
str
input must result inNone
.
Let’s assume the input can only be one of int/float/str
.
Advertisement
Answer
You can use typing.overload
to narrow the types of return values based on the types of the arguments:
from typing import overload @overload def f(value: str) -> None: ... @overload def f(value: int) -> int: ... @overload def f(value: float) -> float: ... def f(value: int|float|str) -> int|float|None: if isinstance(value, str): return None return value
Note that the actual implementation uses a union of all the possible argument and return types.
If a call to the function conforms to one of the overload stubs, the appropriate return type is inferred:
reveal_type(f("foo")) # revealed type is "None" reveal_type(f(123)) # revealed type is "builtins.int"