如何在Python/Mypy中实现关联类型？子类参数类型适配方案

问题背景 #

有一组层级关联的类：BaseProblem/BaseStep的子类对应不同 prover 的具体类型，BaseProver的子类需要重载方法f，接收对应子类类型的参数，但直接重载会触发mypy的类型不兼容错误。用泛型可以解决，但当关联类型数量多（比如10种）时，泛型参数会非常冗长，想要类似Rust关联类型的简洁写法。

原代码示例：

class BaseProblem: ...
class BaseStep: ...

class LeanProblem(BaseProblem): ...
class LeanStep(BaseStep): ...

class IsabelleProblem(BaseProblem): ...
class IsabelleStep(BaseStep): ...

class BaseProver:
    def f(self, problem: BaseProblem, step: BaseStep): ...
class LeanProver(BaseProver):
    def f(self, problem: LeanProblem, step: LeanStep): ...  # mypy报错：参数类型与父类不兼容
class IsabelleProver(BaseProver):
    def f(self, problem: IsabelleProblem, step: IsabelleStep): ...

mypy错误信息：

Argument 1 of "f" is incompatible with supertype "LeanProblem";
supertype defines the argument type as "BaseProblem" [override]

用泛型的解决方案虽然可行，但关联类型越多，泛型参数列表越臃肿：

from typing import TypeVar, Generic

TProblem = TypeVar('TProblem', bound=BaseProblem)
TStep = TypeVar('TStep', bound=BaseStep)
# 若有10种关联类型，需要定义10个TypeVar，BaseProver的泛型参数会非常长

class BaseProver(Generic[TProblem, TStep]):
    def f(self, problem: TProblem, step: TStep): ...
class LeanProver(BaseProver[LeanProblem, LeanStep]):
    def f(self, problem: LeanProblem, step: LeanStep): ...

解决方案：模拟Rust关联类型 #

Python没有原生的关联类型，但可以通过以下几种方式模拟：

1. 类级别类型注解（兼容Python 3.8+） #

在父类协议中定义抽象类属性作为关联类型标记，子类指定具体类型，结合Protocol约束类型一致性：

from typing import TypeVar, Protocol, Type

TBaseProblem = TypeVar('TBaseProblem', bound='BaseProblem')
TBaseStep = TypeVar('TBaseStep', bound='BaseStep')

class BaseProblem: ...
class BaseStep: ...

class ProverProtocol(Protocol[TBaseProblem, TBaseStep]):
    # 定义关联类型：子类必须赋值具体类型
    ProblemType: Type[TBaseProblem]
    StepType: Type[TBaseStep]

    def f(self, problem: TBaseProblem, step: TBaseStep) -> None: ...

class LeanProblem(BaseProblem): ...
class LeanStep(BaseStep): ...

class LeanProver(ProverProtocol[LeanProblem, LeanStep]):
    ProblemType = LeanProblem
    StepType = LeanStep

    def f(self, problem: LeanProblem, step: LeanStep) -> None:
        # 业务逻辑实现
        ...

class IsabelleProblem(BaseProblem): ...
class IsabelleStep(BaseStep): ...

class IsabelleProver(ProverProtocol[IsabelleProblem, IsabelleStep]):
    ProblemType = IsabelleProblem
    StepType = IsabelleStep

    def f(self, problem: IsabelleProblem, step: IsabelleStep) -> None:
        # 业务逻辑实现
        ...

这种方式把关联类型集中在类属性中，即使有10种关联类型，也只需在协议中添加对应抽象属性，子类赋值即可，比冗长的泛型参数更整洁。

2. TypeAlias + 抽象基类（Python 3.10+） #

用TypeAlias定义关联类型，结合抽象基类强制子类覆盖类型别名：

from typing import TypeAlias, ABC, abstractmethod
from abc import ABCMeta

class BaseProblem: ...
class BaseStep: ...

class BaseProver(ABC, metaclass=ABCMeta):
    # 抽象类型别名，子类必须覆盖
    ProblemType: TypeAlias
    StepType: TypeAlias

    @abstractmethod
    def f(self, problem: ProblemType, step: StepType) -> None: ...

class LeanProblem(BaseProblem): ...
class LeanStep(BaseStep): ...

class LeanProver(BaseProver):
    ProblemType = LeanProblem
    StepType = LeanStep

    def f(self, problem: ProblemType, step: StepType) -> None:
        ...

class IsabelleProblem(BaseProblem): ...
class IsabelleStep(BaseStep): ...

class IsabelleProver(BaseProver):
    ProblemType = IsabelleProblem
    StepType = IsabelleStep

    def f(self, problem: ProblemType, step: StepType) -> None:
        ...

mypy会识别子类覆盖的类型别名，自动校验f方法的参数类型，同时保持代码结构简洁。

3. 关联类型打包为容器类（适合大量关联类型场景） #

如果关联类型数量极多，可将所有关联类型打包到单独的容器类中，父类仅需一个泛型参数指向容器：

from typing import TypeVar, Generic, Type

class BaseProblem: ...
class BaseStep: ...
# 假设有更多关联类型
class BaseConfig: ...
class BaseResult: ...

# 定义类型容器基类
class ProverTypes:
    Problem: Type[BaseProblem]
    Step: Type[BaseStep]
    Config: Type[BaseConfig]
    Result: Type[BaseResult]

# Lean对应的类型容器
class LeanTypes(ProverTypes):
    Problem = LeanProblem
    Step = LeanStep
    Config = LeanConfig  # 假设已定义
    Result = LeanResult  # 假设已定义

TProverTypes = TypeVar('TProverTypes', bound=ProverTypes)

class BaseProver(Generic[TProverTypes]):
    def f(self, problem: TProverTypes.Problem, step: TProverTypes.Step) -> None: ...
    def configure(self, config: TProverTypes.Config) -> None: ...
    def get_result(self) -> TProverTypes.Result: ...

class LeanProver(BaseProver[LeanTypes]):
    def f(self, problem: LeanTypes.Problem, step: LeanTypes.Step) -> None: ...
    def configure(self, config: LeanTypes.Config) -> None: ...
    def get_result(self) -> LeanTypes.Result: ...

这种方式将关联类型集中管理，父类泛型参数数量固定，即使扩展10种关联类型，只需更新容器类即可，极大降低维护复杂度。

关键说明 #

• 以上方式均能让mypy正确识别子类方法的参数类型兼容性，避免重载错误。
• 相比原生泛型，模拟关联类型的写法更适合多关联类型场景，结构清晰，维护成本更低。

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