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