how to type a function that returns None if input is string otherwise return the input as is

user_856804 0
Viewed 39

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.

type-annotation mypy python
1 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"
user_37997590
Related