使用 Python 中的類型模塊為具有多個返回類型和條件的函數創建泛型類型提示的正確方法是什么？

Question

在使用pyright工具進行類型檢查時，我遇到了以下問題。

（Python 3.11，版權 1.1.310）

我的代碼：

class TypeCast(str, enum.Enum):
    STRING = str
    FLOAT = float

def cast_int_to_str(value: int) -> str:
    return _cast_int_value(value, type_cast=TypeCast.STRING)


def cast_int_to_float(value: int) -> float:
    return _cast_int_value(value, type_cast=TypeCast.FLOAT)


def _cast_int_value(value: int, type_cast: TypeCast):
    match type_cast:
        case type_cast.STRING:
            return str(value)
        case type_cast.FLOAT:
            return float(value)
        case _:
            raise ValueError(
                f"Invalid type_cast value: {type_cast}"
            )

1. 如果我將_cast_int_value function 的返回類型指定為：

def _cast_int_value(value: int, type_cast: TypeCast) -> str | float:  # typing.Union[str, float] in older Python versions

然后我收到以下 pyright 錯誤：

error: Expression of type "str | float" cannot be assigned to return type "str"
    Type "str | float" cannot be assigned to type "str"
      "float" is incompatible with "str" (reportGeneralTypeIssues)

error: Expression of type "str | float" cannot be assigned to return type "float"
    Type "str | float" cannot be assigned to type "float"
      "str" is incompatible with "float" (reportGeneralTypeIssues)

2. 如果我將 TypeVar 創建為：

from typing import TypeVar
TPossibleCastType = TypeVar("TPossibleCastType", bound=str | float)

並將其指定為_cast_int_value的返回類型 function

def _cast_int_value(value: int, type_cast: TypeCast) -> TPossibleCastType:

版權錯誤是：

error: Expression of type "str" cannot be assigned to return type "TPossibleCastType@_cast_int"
    Type "str" cannot be assigned to type "TPossibleCastType@_cast_int" (reportGeneralTypeIssues)

error: Expression of type "float" cannot be assigned to return type "TPossibleCastType@_cast_int"
    Type "float" cannot be assigned to type "TPossibleCastType@_cast_int" (reportGeneralTypeIssues)

問題：我做錯了什么？ 我應該使用 typing 模塊的什么功能來解決這個問題？

PS 代碼示例是合成的，僅用於說明問題

Answer 1

您可以使用@overload裝飾器來區分類型。

from enum import Enum
from typing import overload, Literal


class MyEnum(Enum):
    FIRST = 1
    SECOND = 2


@overload
def enum_thing(value: int, other: Literal[MyEnum.FIRST]) -> int:
    ...


@overload
def enum_thing(value: int, other: Literal[MyEnum.SECOND]) -> float:
    ...


def enum_thing(value: int, other: MyEnum) -> int | float:
    if other == MyEnum.FIRST:
        return value * 2
    elif other == MyEnum.SECOND:
        return value / 2
    else:  # For linter
        raise ValueError("Unmatched")


a = enum_thing(1, MyEnum.FIRST)  # int
b = enum_thing(1, MyEnum.SECOND)  # float