How to avoid typing errors with functions having a varying return type

[KelSolaar]

Hello,

Given the following code, that uses a function with a varying return type:

from typing import Tuple

class Foo:
    pass

def bar(additional_data: bool = False) -> float | Tuple[float, Foo]:
    if additional_data:
        return 1, Foo()
    else:
        return 1

a = bar(False)
b, c = bar(True)

Pyright produces the following error:

13:8 -  error: "float" is not iterable
    "__iter__" method not defined (reportGeneralTypeIssues)

I know I can do a check on the return value as follows to prevent the error:

a = bar(False)
result = bar(True)
if isinstance(result, tuple):
    b, c = result

I was wondering however if there is another way?

[erictraut]

The Python type system includes the concept of an "overload", which allows you to specify polymorphic function behaviors where a function's return type depends on the types of its parameters.

Here's how that would be done with your example:

from typing import Literal, overload

class Foo:
    pass

@overload
def bar(additional_data: Literal[True]) -> tuple[float, Foo]: ...
@overload
def bar(additional_data: Literal[False] = False) -> float: ...
def bar(additional_data: bool = False) -> float | tuple[float, Foo]:
    if additional_data:
        return 1, Foo()
    else:
        return 1

a = bar(False)
b, c = bar(True)

For more details, refer to the Python typing spec.

[KelSolaar]

Ah of course, beautiful! Totally forgotten about this as I never had the use case, until now :)

Thank you so much!

[KelSolaar]

Turns out that my code is more complicated and the additional_data arg is optional:

from typing import Literal, overload


@overload
def bar(
    baz: str = "Bob", additional_data: Literal[True] = True
) -> tuple[float, float, float]: ...


@overload
def bar(
    baz: str = "Bob", additional_data: Literal[False] = False
) -> tuple[float, float]: ...


def bar(
    baz: str = "Bob", additional_data: bool = False
) -> tuple[float, float] | tuple[float, float, float]:
    if additional_data:
        return 1, 2, 3
    else:
        return 1, 2


a, c = bar("John", False)
d, e, f = bar("John", True)

So I get this one:

5:5 - error: Overload 1 for "bar" overlaps overload 2 and returns an incompatible type (reportOverlappingOverload)

[KelSolaar]

from typing import Literal, overload

@overload
def bar(
baz: str = ..., additional_data: Literal[True] = True
) -> tuple[float, float, float]: ...

@overload
def bar(
baz: str = ..., *, additional_data: Literal[False]
) -> tuple[float, float]: ...

@overload
def bar(baz: str, additional_data: Literal[False]) -> tuple[float, float]: ...

def bar(
baz: str = "Bob", additional_data: bool = False
) -> tuple[float, float] | tuple[float, float, float]:
if additional_data:
return 1, 2, 3
else:
return 1, 2

a, b = bar("John", False)
c, d, e= bar("John", True)