Python typing, if none is given return none

Question

Is it possible to write a typehint in python that guarantees if None is given to a function then None is returned?

For example this is possible:

from typing import Dict, Union
def getMaybe(dictionary: Optional[Dict], key: str) -> Optional[str]:
    if dictionary is None:
        return dictionary

    return dictionary.get(key)

But then even when I know the argument has a value, the type signature doesn't guarantee the output will. For example:

def printer(msg: str):
    print(msg)

data = {'a': 'a'}
result = getMaybe(data, 'a')
printer(result)

Gives the error:

error: Argument of type "str | None" cannot be assigned to parameter "msg" of type "str" in function "printer"
 ?Type "str | None" cannot be assigned to type "str"
  ??Type "None" cannot be assigned to type "str" (reportGeneralTypeIssues)

Is it possible to encode in the type signature that when None is given as an argument, then None is returned?

question from:https://stackoverflow.com/questions/65909744/python-typing-if-none-is-given-return-none

Answer 1

typing.overload is what you want:

from typing import Dict, Union, Optional, overload

@overload
def getMaybe(dictionary: None, key: str) -> None: ...

@overload
def getMaybe(dictionary: Dict, key: str) -> str: ...
    
def getMaybe(dictionary: Optional[Dict], key: str) -> Optional[str]:
    if dictionary is None:
        return dictionary

    return dictionary.get(key)
    

reveal_type(getMaybe(None, "")) # None
reveal_type(getMaybe({}, "")) # str