Correctly typing a function that can return a provided default value

user_235217 0
Viewed 144

I have a function of the following structure:

def get_something_from_data(data: Mapping[str, str], default: Optional[str] = None) -> Optional[str]:
    """Get something out of `data` if it is there, if not return the value of `default`. 
    If `default` is not provided, return None.
    """

So this can return either Optional[str] if default is omitted, or will always return str if default has an str value.

Now I call this in code similar to:

has_value = get_something_from_data(data, "fallback")
return has_value.endswith("k")

Even though in this case has_value will always be an str, running mypy on this code generates an error because it thinks has_value might be None and will not have endswith.

So I tried typing my function differently:

DT = TypeVar('DT', str, None)

def get_something_from_data(data: Mapping[str, str], default: DT = None) -> Union[str, DT]:
    pass

This works, but now when I call mypy on code like:

has_value = get_something_from_data(data)
return has_value.endswith("k")  # might fail as has_value can be None

mypy doesn't throw an error, even though the code is potentially dangerous.

Is there a way to properly type this function so mypy generates errors based on whether or not a default value was provided?

python-typing mypy type-hinting python
1 Answers

Using typing.overload you can describe multiple combinations of arguments and return types for a function

from typing import overload, Mapping, Optional


@overload
def get_something_from_data(data: Mapping[str, str], default: None = None) -> Optional[str]:
    ...


@overload
def get_something_from_data(data: Mapping[str, str], default: str) -> str:
    ...


def get_something_from_data(data, default=None):
    pass
user_5485620
Related