如何對返回類型取決於參數的輸入類型的函數進行類型提示？

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

這正是函數重載的用途。

簡而言之，您執行以下操作：

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 。

在運行時，實際運行的唯一功能是最后一個 - 前兩個被完全覆蓋並被忽略。

Answer 2

如果您的返回類型與您的輸入類型相同（或者，如我的示例所示，一個變體），您可以使用以下策略（並且在引入/支持更多類型時無需添加額外的@overload ）。

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)

這個類型定義比使用這個更具體

def overly_general_filter(category: Any) -> list[Any]:
  pass

它真的相當於

def equally_general_filter(category) -> list:
  pass