使用 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