如何为参数类型互相关联的函数添加类型注解且避免代码重复？

问题背景 #

给类的__init__方法添加类型注解时，前两个参数x和y存在类型关联（比如x:A时必须对应y:B，x:C时必须对应y:D），但方法还有20多个与这两个参数无关的其他参数。直接使用@overload会因重载签名缺少其他参数导致mypy报错，重复所有参数又会产生大量冗余代码，违反DRY原则。

解决方案 #

方案一：用TypedDict复用公共关键字参数 #

定义一个包含所有公共关键字参数的TypedDict，在重载签名中通过**kwargs引用，避免重复定义参数。

from __future__ import annotations
from typing import overload, TypedDict
from typing_extensions import TypeAlias

A: TypeAlias = str
B: TypeAlias = int
C: TypeAlias = list
D: TypeAlias = tuple

# 定义所有公共关键字参数的类型约束
class ExampleInitKwargs(TypedDict, total=False):
    foo: str
    bar: str
    option: str
    # ... 补充剩余20个可选参数 ...

class Example:
    @overload
    def __init__(self, x: A, y: B, **kwargs: ExampleInitKwargs): ...

    @overload
    def __init__(self, x: C, y: D, **kwargs: ExampleInitKwargs): ...

    def __init__(
        self,
        x,
        y,
        *,
        foo: str,
        bar: str,
        option: str = '',
        # ... 20个可选参数 ...
    ):
        pass  # 业务逻辑

说明：该方案完全兼容原代码的参数传递方式（公共参数需以关键字形式传递），既保留了x和y的类型关联检查，又避免了代码冗余。

方案二：用Protocol定义初始化签名 #

通过Protocol分别定义两种参数组合的完整签名，再将类的__init__注解为这些协议的联合类型。

from __future__ import annotations
from typing import overload, Protocol
from typing_extensions import TypeAlias

A: TypeAlias = str
B: TypeAlias = int
C: TypeAlias = list
D: TypeAlias = tuple

# 定义第一种参数组合的初始化协议
class ExampleInitAB(Protocol):
    def __call__(self, x: A, y: B, *, foo: str, bar: str, option: str = '', ...) -> None: ...

# 定义第二种参数组合的初始化协议
class ExampleInitCD(Protocol):
    def __call__(self, x: C, y: D, *, foo: str, bar: str, option: str = '', ...) -> None: ...

class Example:
    __init__: ExampleInitAB | ExampleInitCD

    def __init__(
        self,
        x,
        y,
        *,
        foo: str,
        bar: str,
        option: str = '',
        # ... 20个可选参数 ...
    ):
        pass  # 业务逻辑

说明：Protocol中的...表示允许额外参数，mypy会根据协议进行严格的类型检查，实现部分只需维护一份参数列表。

方案三：用泛型类关联参数类型 #

通过TypeVar定义关联的类型变量，结合Generic类约束x和y的合法类型组合，无需使用@overload。

from __future__ import annotations
from typing import TypeVar, Generic
from typing_extensions import TypeAlias

A: TypeAlias = str
B: TypeAlias = int
C: TypeAlias = list
D: TypeAlias = tuple

# 定义关联的类型变量，限定允许的类型组合
X = TypeVar('X', A, C)
Y = TypeVar('Y', B, D)

class Example(Generic[X, Y]):
    def __init__(
        self,
        x: X,
        y: Y,
        *,
        foo: str,
        bar: str,
        option: str = '',
        # ... 20个可选参数 ...
    ):
        pass  # 业务逻辑

说明：实例化时如果类型推断不明确，可显式指定类型参数，比如Example[A, B]("test", 123, foo="a", bar="b")。这种方式适合参数关联关系更通用的场景。

方案选择建议 #

• 若公共参数均为关键字参数，优先选方案一，简洁且兼容性好。
• 若需要严格的签名校验，选方案二。
• 若参数关联是通用的类型映射关系，选方案三的泛型实现。

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