如何对返回类型取决于参数的输入类型的函数进行类型提示？

Posted 2023-03-22

【中文标题】如何对返回类型取决于参数的输入类型的函数进行类型提示？

【英文标题】：How can I type-hint a function where the return type depends on the input type of an argument?

【发布时间】：2019-02-25 23:19:19

【问题描述】：

假设我有一个函数可以将 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)

我知道我可以将映射拆分为两个函数，但目的是拥有一个通用类型映射函数。

我也在考虑使用类和多态性，但是我将如何对最顶层的类方法进行类型提示呢？因为它们的输出类型取决于具体的实例类型。

【问题讨论】：

一个问题是：大概你在运行时不知道输入值是什么，所以你也不知道输出是什么。所以你必须期望这两种类型都是函数的返回值。

是的。但正如我在删除的答案中提到的，我认为这可能是singledispatch 的工作。但我现在将问题悬而未决。因为如果它真的不可行，“你不能”将是一个很好的答案。

你为什么不通过签名让它多态呢？ （不知道这在 python 中是如何工作的，但在 scala 中你一直这样做），或者在内部定义匿名函数并根据你得到的返回它们的类型。非常有趣的问题顺便说一句

@E.Serra 你有指向 Scala 示例的链接吗？也许我可以把它当作灵感，或者至少作为我学习新术语的来源，我可以用它来改进我的谷歌搜索；）

【参考方案1】：

这正是function overloads 的用途。

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

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。

在运行时，唯一实际运行的函数是最后一个——前两个被完全覆盖和忽略。

【讨论】：

是的，这就是我的意思，你一直在 scala 中这样做，你甚至不必指定重载，因为函数绑定到它们的签名。