如何键入提示函数,其中返回类型取决于参数的输入类型？

Question

假设我有一个将Python数据类型转换为Postgres数据类型的函数,如下所示:

def map_type(input):
    if isinstance(input, int):
        return MyEnum(input)
    elif isinstance(input, str):
        return MyCustomClass(str)

我可以输入提示:

def map_type(input: Union[int, str]) -> Union[MyEnum, MyCustomClass]: ...

但是,即使它是正确的,下面的代码也无法进行类型检查:

myvar = map_type('foobar')
print(myvar.property_of_my_custom_class)

完整示例(工作代码,但类型提示中的错误):

from typing import Union
from enum import Enum


class MyEnum(Enum):
    VALUE_1 = 1
    VALUE_2 = 2


class MyCustomClass:

    def __init__(self, value: str) -> None:
        self.value = value

    @property
    def myproperty(self) -> str:
        return 2 * self.value


def map_type(value: Union[int, str]) -> Union[MyEnum, MyCustomClass]:

    if isinstance(value, int):
        return MyEnum(value)
    elif isinstance(value, str):
        return MyCustomClass(value)
    raise TypeError('Invalid input type')


myvar1 = map_type(1)
print(myvar1.value, myvar1.name)

myvar2 = map_type('foobar')
print(myvar2.myproperty)

我知道我可以将映射拆分为两个函数,但目标是具有泛型类型映射函数.

我也在考虑使用类和多态,但是我如何键入提示最顶级的类方法呢？因为它们的输出类型取决于具体的实例类型.

Answer 1

对于相同的输出类型或输入类型的变体

@overload您可以使用以下策略（并且当引入/支持更多类型时无需添加额外的s）。

from typing import TypeVar

T = TypeVar("T")  # the variable name must coincide with the string

def filter_category(category: T) -> list[T]:
  # assume get_options() is some function that gets
  # an arbitrary number of objects of varying types
  return [
    option
    for option in get_options()
    if is subclass(option, category)
  ]

filter_category将由您的 IDE（VSCode、Pycharm 等）和静态类型检查器（Mypy）正确关联

# for instance list_of_ints would show in your IDE as of type list[Type[int]]
list_of_ints = filter_category(int)

# or here it list_of_bools would show in your IDE as of type list[Type[bool]]
list_of_bools = filter_category(bool)

注意：这并不能回答问题的每个方面，只能回答其中的一个子集，尽管这是一个非常常见的问题

Answer 2

这正是函数重载的目的。

简而言之，您可以执行以下操作：

from typing import overload

# ...snip...

@overload
def map_type(value: int) -> MyEnum: ...

@overload
def map_type(value: str) -> MyCustomClass: ...

def map_type(value: Union[int, str]) -> Union[MyEnum, MyCustomClass]:
    if isinstance(value, int):
        return MyEnum(value)
    elif isinstance(value, str):
        return MyCustomClass(value)
    raise TypeError('Invalid input type')

现在，当您这样做时map_type(3)，mypy将了解返回类型为MyEnum。

在运行时，唯一要实际运行的功能是最后一个功能-前两个功能完全被覆盖和忽略。