[Python] How to annotate python function with different return type based on parameters?

There are some cases when a function will return different type based on the parameter input.

A simple example could be

from typing import Union

def func(is_str: bool) -> Union[str, int]:
if is_str:
return "string"
else:
return 42


a:int = func(is_str=False)


For pylint, it will raise a warning:

Type "str | int" is not assignable to declared type "int".


I wonder what is the best practice for circumstances like this.

What I tried


One suggested way is to rewrite the function with decorator @overload, but this scenario also happens for some package dependencies, when the function body is relative long (e.g. a function handling network request, and decide whether to return a json type(response.json()) or a string(response.text).)


Another way for this is

from typing import TypeVar

T = TypeVar('T')

def func(return_type:T = str) -> T:
if return_type is str:
return "string"
else:
return 42


That would also raise some warnings.

Type "Literal[string]" is not assignable to type "T@func"
Type "Literal[42]" is not assignable to type "T@func"


Is is possible to avoid the warning without modification the function?


Another approach for solving this would be:

from typing import Union

def func(is_str: bool) -> Union[str, int]:
if is_str:
return "string"
else:
return 42

a:Union[str, int] = func(is_str=False)

if isinstance(a, str):
raise TypeError('parameter a should be an integer.')


So that variable a will be considered an integer for pylint in later code.

this practice is needed when the code is long, but for short code cases, it is just obvious, and write the checking can be tiring especially when calling the function multiple times.

Continue reading...