Let's say that I have a function returning a single value in some cases and multiple values (i.e. tuple) in other cases.
def foo(a: int, b: int, flag: bool):
if flag:
return a
else:
return a, b
What is the correct type annotation for the return value of this fuction?
I can propose something like this:
def foo(a: int, b: int, flag: bool) -> Tuple[int, Optional[int]]:
if flag:
return a
else:
return a, b
But the main thing that I don't like in this code is that in case when the function returns a single value, that value is not a tuple. So to make this work properly I should remake the function to
def foo(a: int, b: int, flag: bool) -> Tuple[int, Optional[int]]:
if flag:
return (a,)
else:
return a, b
Which is something that I don't want to do. Any thoughts?
You want to use typing.Union
to indicate that the function can return either one type of thing or another. For example, if you had a function that could either return an int
or a str
, then it's return type would be Union[int, str]
.
So,
def foo(a: int, b: int, flag: bool) -> Union[int, Tuple[int, int]]:
If you're using Python 3.10, there's a better way to write a union: int | Tuple[int, int]
However, I should repeat the warning given in the comments. Having such a return type is an antipattern and you should refactor your code to avoid this.