Python 3 base class type annotation only allow the current subclass

Question

Let's say I've got three classes, one parent class, and two subclasses:

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

class ChildA(BaseModel):
    pass

class ChildB(BaseModel):
    pass

The parent class has a method that takes another instance of the current class and returns a new instance of the current class (out of scope for this question).

How do I annotate BaseModel.merge to restrict it to only the current subclass?

I can do something like this:

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

But this still allows me to pass an instance of ChildB into ChildA, since both inherit from BaseModel. I only want ChildA to be allowed in ChildA, and ChildB to be allowed for ChildB. How can I do that without reimplementing merge on each subclass?

Answer 1

Annotate both arguments with a type variable, to enforce that both arguments must be of the same type.

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