如何为类型到处理函数的映射字典添加正确的Python类型提示？

我需要创建一个字典，将输入数据类型的元组映射到对应处理该类型组合的函数，但无法写出能通过mypy检查的类型提示。以下是最小可复现代码：

from typing import Callable

class BaseType: pass
class IntType(BaseType): pass
class StringType(BaseType): pass

def _compare_int_int(before: IntType, after: IntType) -> int:
    return 1

def _compare_int_str(before: IntType, after: StringType) -> int:
    return 2

def _compare_str_int(before: StringType, after: IntType) -> int:
    return 3

def _compare_str_str(before: StringType, after: StringType) -> int:
    return 4

_COMPARERS: dict[tuple[type[BaseType], type[BaseType]], Callable[[BaseType, BaseType], int]] = {
    (IntType, IntType): _compare_int_int,
    (IntType, StringType): _compare_int_str,
    (StringType, StringType): _compare_str_str,
    (StringType, IntType): _compare_str_int
}

def compare(before: BaseType, after: BaseType) -> int:
    compare_func = _COMPARERS[(type(before), type(after))]
    return compare_func(before, after)

a = IntType()
b = StringType()

print(compare(a, a))  # 1
print(compare(a, b))  # 2
print(compare(b, a))  # 3
print(compare(b, b))  # 4

代码运行逻辑正常，但mypy抛出如下错误：

Dict entry 0 has incompatible type "Tuple[Type[IntType], Type[IntType]]": "Callable[[IntType, IntType], int]"; expected "Tuple[Type[BaseType], Type[BaseType]]": "Callable[[BaseType, BaseType], int]"

其余字典条目也会触发类似的类型不兼容错误。

我尝试过使用TypeVar来约束类型：

T = TypeVar('T', bound=BaseType)
U = TypeVar('U', bound=BaseType)

_COMPARERS: dict[tuple[type[T], type[U]], Callable[[T, U], int]] = {...}

这个写法逻辑更清晰，但仍然触发相同类型错误，求正确的类型提示写法。

解决方案 #

方法一：使用Protocol定义兼容签名 + 类型断言 #

mypy无法自动推导字典中键与值的类型对应关系，我们可以先定义一个兼容所有子类型的函数协议，再通过cast辅助类型检查：

from typing import Callable, TypeVar, cast, Protocol

class BaseType: pass
class IntType(BaseType): pass
class StringType(BaseType): pass

# 定义协议，允许接受BaseType的子类实例作为参数
class Comparer(Protocol):
    def __call__(self, before: BaseType, after: BaseType) -> int: ...

def _compare_int_int(before: IntType, after: IntType) -> int:
    return 1

def _compare_int_str(before: IntType, after: StringType) -> int:
    return 2

def _compare_str_int(before: StringType, after: IntType) -> int:
    return 3

def _compare_str_str(before: StringType, after: StringType) -> int:
    return 4

# 用协议作为字典值的类型，让所有子类型函数都能兼容
_COMPARERS: dict[tuple[type[BaseType], type[BaseType]], Comparer] = {
    (IntType, IntType): _compare_int_int,
    (IntType, StringType): _compare_int_str,
    (StringType, StringType): _compare_str_str,
    (StringType, IntType): _compare_str_int
}

T = TypeVar('T', bound=BaseType)
U = TypeVar('U', bound=BaseType)

def compare(before: T, after: U) -> int:
    # 通过cast告诉mypy，当前获取的函数可以处理T和U类型
    compare_func = cast(Callable[[T, U], int], _COMPARERS[(type(before), type(after))])
    return compare_func(before, after)

方法二：使用TypedDict枚举所有固定键值对（适用于类型组合有限的场景） #

如果你的类型组合是固定的，可以用TypedDict明确列出每个键对应的精确函数类型，让mypy进行精确推导：

from typing import Callable, TypedDict, cast

class BaseType: pass
class IntType(BaseType): pass
class StringType(BaseType): pass

def _compare_int_int(before: IntType, after: IntType) -> int:
    return 1

def _compare_int_str(before: IntType, after: StringType) -> int:
    return 2

def _compare_str_int(before: StringType, after: IntType) -> int:
    return 3

def _compare_str_str(before: StringType, after: StringType) -> int:
    return 4

# 用TypedDict定义每个键对应的精确函数类型
class ComparerDict(TypedDict):
    (IntType, IntType): Callable[[IntType, IntType], int]
    (IntType, StringType): Callable[[IntType, StringType], int]
    (StringType, IntType): Callable[[StringType, IntType], int]
    (StringType, StringType): Callable[[StringType, StringType], int]

_COMPARERS: ComparerDict = {
    (IntType, IntType): _compare_int_int,
    (IntType, StringType): _compare_int_str,
    (StringType, IntType): _compare_str_int,
    (StringType, StringType): _compare_str_str
}

def compare(before: BaseType, after: BaseType) -> int:
    key = (type(before), type(after))
    assert key in _COMPARERS, f"No comparer found for types {key}"
    compare_func = _COMPARERS[key]
    # 通过cast确保函数参数与输入类型兼容
    return cast(Callable[[BaseType, BaseType], int], compare_func)(before, after)

原写法失效的原因 #

mypy对字典的类型检查是同构的：要求所有键都符合统一的键类型，所有值都符合统一的值类型。而Callable[[IntType, IntType], int]是Callable[[BaseType, BaseType], int]的逆变类型——函数参数更具体，无法被视为基类型函数的子类型，因此mypy判定类型不兼容。

内容的提问来源于stack exchange，提问作者TCF