pythontype-hinting

Correct type hint for multiple return values


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?


Solution

  • 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.