Перегрузка функции с аргументами по умолчанию

Question

У меня есть функция, которая в зависимости от bool-аргумента возвращает тот или иной тип, наподобие такой:

from typing import Union
def foo(num: int = 1, as_bytes: bool = False) -> Union[str, bytes]:
    rv = str(num)
    return rv.encode("utf-8") if as_bytes else rv

foo(1, False)  # Revealed type is 'Union[builtins.str, builtins.bytes]'
foo(as_bytes=True)  # Revealed type is 'Union[builtins.str, builtins.bytes]'

Я хочу написать аннотации для перегрузки функции так, чтобы mypy понимал, какой именно тип возвращается в зависимости от значения аргумента:

foo(1, False)  # Revealed type is 'builtins.str'
foo(as_bytes=True)  # Revealed type is 'builtins.bytes'

Самое лучшее, что у меня получилось написать — это вот такое:

from typing import Union, Literal, overload
@overload
def foo(num: int = 1, as_bytes: Literal[False] = False) -> str: ...
@overload
def foo(num: int = 1, as_bytes: Literal[True] = True) -> bytes: ...
@overload
def foo(num: int = 1, as_bytes: bool = False) -> Union[str, bytes]: ...

Но проблема в том, что mypy в таком случае ругается ошибкой «Overloaded function signatures 1 and 2 overlap with incompatible return types». И это вполне логично, ведь вызов foo() будет соответствовать двум разным сигнатурам. Однако я так и не придумал, как это исправить. Если попытаться написать as_bytes: Literal[True] = False, то ругается ошибкой «Incompatible default for argument», что тоже логично.

Как правильно написать эти аннотации, если это вообще возможно?

Answer 1

Получился такой вариант

from typing import Union, Literal, overload


@overload
def foo(num: int = 1) -> str:  ...
@overload
def foo(*, as_bytes: Literal[True]) -> bytes:  ...
@overload
def foo(*, as_bytes: Literal[False]) -> str:  ...
@overload
def foo(num: int, as_bytes: Literal[True]) -> bytes: ...
@overload
def foo(num: int, as_bytes: Literal[False]) -> str: ...


def foo(num: int = 1, as_bytes: bool = False) -> Union[str, bytes]:
    rv = str(num)
    return rv.encode("utf-8") if as_bytes else rv



s = foo(1, False)
b = foo(as_bytes=True)
s2 = foo(1)

вроде не выдает ошибок и правильно определяет все типы

Answer 2

Перегруженные функции должны отличаться типами аргументов, а не значениями аргументов и не типом возвращаемого значения. Он ругается на это, а не на аннотации.

БЛОГ НА HUSL

Перегрузка функции с аргументами по умолчанию

Ответы (2 шт):