Python3基类类型注释只允许当前的子类

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