Python 3 基类类型注解只允许当前子类

Question

假设我有三个类，一个父类和两个子类:

class BaseModel:
    def merge(self, other):
        return self + other

class ChildA(BaseModel):
    pass

class ChildB(BaseModel):
    pass

父类有一个方法接受当前类的另一个实例并返回当前类的一个新实例(超出这个问题的范围)。

如何注释 BaseModel.merge 以将其限制为仅当前子类？

我可以这样做:

def merge(self, other: BaseModel) -> BaseModel:
    return self + other

但这仍然允许我将 ChildB 的实例传递给 ChildA，因为它们都继承自 BaseModel。我只希望在 ChildA 中允许使用 ChildA，在 ChildB 中允许使用 ChildB。如果不在每个子类上重新实现 merge，我该如何做到这一点？

Answer 1

用类型变量注释两个参数，以强制两个参数必须是同一类型。

from typing import TypeVar

B = TypeVar('B', bound='BaseModel')

class BaseModel:
    def __init__(self, x: int):
        self.x = x

    def __add__(self: B, other: B) -> B:
        return type(self)(self.x + other.x)

    def merge(self: B, other: B) -> B:
        return self + other

class ChildA(BaseModel):
    pass

class ChildB(BaseModel):
    pass


print(ChildA(3).merge(ChildA(4)).x)  # Valid; both arguments are ChildA      
print(ChildA(3).merge(ChildB(4)).x)  # Invalid; one ChildA and one ChildB