Есть ли способ автоматически перегружать функции с соответствием типов ввода-вывода 1:1? ⇐ Python

[Anonymous]

Предположим, у меня есть эта функция

Код: Выделить всё

from typing_extensions import reveal_type

def func(x: str | list[str]) -> int | list[int]:
"""
Some docstring that should show up.
"""
if isinstance(x, str):
return 0
return list(range(len(x)))

y = func(["hi", "hello"])
reveal_type(y)  # int | list[int]

Показ типа y или наведение на него курсора в вашей IDE сообщит вам int | list[int].
Чтобы сделать тип y более точным, мы перегружаем его:

Код: Выделить всё

from typing import overload

from typing_extensions import reveal_type

@overload
def func(x: str) -> int: ...

@overload
def func(x: list[str]) -> list[int]: ...

def func(x: str | list[str]) -> int | list[int]:
"""
Some docstring that should show up.
"""
if isinstance(x, str):
return 0
return list(range(len(x)))

y = func(["hi", "hello"])
reveal_type(y)  # list[int]

Теперь тип y — list[int].
У меня много таких функции — те, в которых возможные типы возвращаемых значений имеют соответствие 1-к-1 с типами одного входа функции. В идеале мне просто нужно было бы сделать что-то вроде—

Код: Выделить всё

from typing import UnionMatch  # not a real thing

def func(x: UnionMatch[str, list[str]]) -> UnionMatch[int, list[int]]:
"""
Some docstring that should show up.
"""
if isinstance(x, str):
return 0
return list(range(len(x)))

— точнее, о типе возвращаемого значения. Никаких @overload не требуется.
Как можно достичь такой функциональности?
Пользователь kevdog824 на Reddit предлагает решение здесь и здесь для случая, когда в UnionMatch ровно 2 типа:

Код: Выделить всё

from typing import Callable, Concatenate, Generic, overload, ParamSpec, TypeVar

from typing_extensions import reveal_type

_C1 = TypeVar("_C1")
_C2 = TypeVar("_C2")
_P = ParamSpec("_P")
_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")

class FuncTemplate(Generic[_C1, _C2, _P, _T1, _T2]):

def __init__(self, func: Callable[Concatenate[_C1 | _C2, _P], _T1 | _T2]) -> None:
self.func = func
# self.__call__.__doc__ = func.__doc__
# Not correct b/c:
# AttributeError: attribute '__doc__' of 'method' objects is not writable

@overload
def __call__(
self,
common_input: _C1,
*args: _P.args,
**kwargs: _P.kwargs,
) -> _T1: ...

@overload
def __call__(
self,
common_input: _C2,
*args: _P.args,
**kwargs: _P.kwargs,
) -> _T2: ...

def __call__(
self,
common_input: _C1 | _C2,
*args: _P.args,
**kwargs: _P.kwargs,
) -> _T1 | _T2:
return self.func(common_input, *args, **kwargs)

@FuncTemplate
def func(x: str | list[str]) -> int | list[int]:
"""
Some docstring that should show up.
"""
if isinstance(x, str):
return 0
return list(range(len(x)))

y = func(["hi", "hello"])
reveal_type(y)  # list[int]

z = func("hey")
reveal_type(z)  # int

Нарушителем соглашения является то, что строка документации функции потеряна. При наведении курсора на функцию в моей IDE (я думаю, она использует Pylance) появляется сообщение

Код: Выделить всё

(function) func: FuncTemplate[str, list[str], (), int, list[int]]

Мне не удалось заставить @functools.wraps работать здесь. И в идеале есть решение, которое работает более чем с двумя типами в UnionMatch.

Подробнее здесь: https://stackoverflow.com/questions/786 ... -output-ty