how to type a function that returns None if input is string otherwise return the input as is
Question:
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 an int
output.
- a
float
input must result in a float
output.
- a
str
input must result in None
.
Let’s assume the input can only be one of int/float/str
.
Answers:
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"
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
.
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"