CODEWITHSUNDEEP
pythonpython-typing
Jess
Jess

Reputation: 8700

Python typing, if none is given return none

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?

Upvotes: 2

Views: 486

Answers (1)

alex_noname
alex_noname

Reputation: 32253

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

Upvotes: 5

Related Questions