使用策略模式时消除mypy的arg-type类型错误

我用策略模式处理不同版本的事件，代码如下：

from typing import overload, TypeVar, Generic

class EventV1: pass
class EventV2: pass

class DataGathererV1:
    def process(self, event: EventV1): pass
    def process2(self, event: EventV1): pass

class DataGathererV2:
    def process(self, event: EventV2): pass
    def process2(self, event: EventV2): pass


class Dispatcher:
    def __init__(self):
        self.worker_v1: DataGathererV1 = DataGathererV1()
        self.worker_v2: DataGathererV2 = DataGathererV2()

    def dispatch(self, event: EventV1 | EventV2):
        handler: DataGathererV1 | DataGathererV2 = self.worker_v1 if isinstance(event, EventV1) else self.worker_v2

        handler.process(event)
        # 公共逻辑

        handler.process2(event)
        # 公共逻辑

        handler.process(event)
        # 其他逻辑...

但mypy报出了一系列arg-type错误：

example.py:36: error: Argument 1 to "process" of "DataGathererV1" has incompatible type "EventV1 | EventV2"; expected "EventV1"  [arg-type]
example.py:36: error: Argument 1 to "process" of "DataGathererV2" has incompatible type "EventV1 | EventV2"; expected "EventV2"  [arg-type]
example.py:40: error: Argument 1 to "process2" of "DataGathererV1" has incompatible type "EventV1 | EventV2"; expected "EventV1"  [arg-type]
example.py:40: error: Argument 1 to "process2" of "DataGathererV2" has incompatible type "EventV1 | EventV2"; expected "EventV2"  [arg-type]
example.py:44: error: Argument 1 to "process" of "DataGathererV1" has incompatible type "EventV1 | EventV2"; expected "EventV1"  [arg-type]
example.py:44: error: Argument 1 to "process" of "DataGathererV2" has incompatible type "EventV1 | EventV2"; expected "EventV2"  [arg-type]

我尝试过重载私有函数来关联处理器与事件类型，但问题依旧：

@overload
def __handler(self, event: EventV1) -> tuple[DataGathererV1, EventV1]:
    ...

@overload
def __handler(self, event: EventV2) -> tuple[DataGathererV2, EventV2]:
    ...

def __handler(self, event: EventV1 | EventV2) -> tuple[DataGathererV1 | DataGathererV2, EventV1 | EventV2]:
    return (
        self.worker_v1 if isinstance(event, EventV1) else self.worker_v2,
        event
    )

我的需求是：让mypy正确识别event和handler的类型关联，同时不修改DataGathererVX的事件参数类型，也不添加断言，保留错误类型事件的触发校验。

解决方案 #

使用Python 3.10+引入的match-case语句，mypy能够通过模式匹配进行类型窄化，自动识别事件与处理器的对应关系，无需拆分dispatch方法或添加额外断言：

from typing import overload, TypeVar, Generic

class EventV1: pass
class EventV2: pass

class DataGathererV1:
    def process(self, event: EventV1): pass
    def process2(self, event: EventV1): pass

class DataGathererV2:
    def process(self, event: EventV2): pass
    def process2(self, event: EventV2): pass


class Dispatcher:
    def __init__(self):
        self.worker_v1: DataGathererV1 = DataGathererV1()
        self.worker_v2: DataGathererV2 = DataGathererV2()

    def dispatch(self, event: EventV1 | EventV2):
        # 使用match-case进行类型匹配，mypy会自动窄化handler的类型
        match event:
            case EventV1():
                handler = self.worker_v1
            case EventV2():
                handler = self.worker_v2

        handler.process(event)
        # 公共逻辑

        handler.process2(event)
        # 公共逻辑

        handler.process(event)
        # 其他逻辑...

原理说明 #

match-case语句会触发mypy的类型守卫机制：当匹配到EventV1()时，mypy会推断出event的类型是EventV1，同时handler的类型是DataGathererV1；匹配到EventV2()时同理。这样后续调用handler.process(event)时，mypy能确认参数类型完全匹配，不会再报类型错误。

这种方案既保留了原有代码的结构（无需拆分dispatch），也没有修改DataGathererVX的类型定义，同时保留了错误类型事件的校验能力——如果传入不符合要求的事件类型，mypy会在调用dispatch时直接报错。

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