I am struggling to annotate types for some classes and subclasses.
Here is the situation before annotations. The problem at hand is how to annotate getOther() properly so that static type checkers correctly infer the return type for all sub-classes, without having to redefine it for the subcalsses SubB1 and SubB2.
Also, SubB1 and SubB2 must remain subclasses of B.
class A:
def __init__(self, a:int) -> None:
self.a:int = a
class SubA1(A): ...
class SubA2(A): ...
class B:
_otherClass = A
def __init__(self, b:int) -> None:
self.b:int = b
def getOther(self): # <<< This is the method I want to annotate
return self._otherClass(a = self.b)
class SubB1(B):
_otherClass = SubA1
class SubB2(B):
_otherClass = SubA2
This is the desired behaviour:
b0 = B(0)
b1 = SubB1(1)
b2 = SubB2(2)
b0.getOther() # Returns A(0) => type checker: A
b1.getOther() # Returns SubA1(1) => type checker: SubA1
b2.getOther() # Returns SubA2(2) => type checker: SubA2
I have tried multiple approaches, but each failed for one reason or another:
1. Making B a generic class:
from typing import TypeVar, Generic
class A:
def __init__(self, a:int) -> None:
self.a:int = a
class SubA1(A): ...
class SubA2(A): ...
Aclass = TypeVar('Aclass')
class B(Generic[Aclass]):
_otherClass = A
def __init__(self, b:int) -> None:
self.b:int = b
def getOther(self) -> Aclass:
return self._otherClass(a = self.b)
class SubB1(B[SubA1]):
_otherClass = SubA1
class SubB2(B[SubA2]):
_otherClass = SubA2
This works for SubB1 and SubB2, but not for B itself, since nowhere is specified that B should work with A as a concrete type (B is actually still a generic class):
b0 = B(0)
b1 = SubB1(1)
b2 = SubB2(2)
_ = b0.getOther() # Type checker => Any, Not OK
_ = b1.getOther() # Type checker => SubA1, OK
_ = b2.getOther() # Type checker => SubA2, OK
2. Separate generic class for getOther
The second attempt was to create a separate generic class to hold the getOther method, and have the other classes inherit from it using multiple inheritance:
...
class _GenericB(Generic[Aclass]):
def getOther(self) -> Aclass:
return self._otherClass(a = self.b)
class B(_GenericB[A]):
_otherClass = A
def __init__(self, b:int) -> None:
self.b:int = b
class SubB1(B, _GenericB[SubA1])
class SubB1(B, _GenericB[SubA2])
Now, this works for B, but not for SubB1 and SubB2, as multiple inheritance is such that getOther called on SubB1 and SubB2 resolves to the method defined in B, which is annotated with A.
b0 = B(0)
b1 = SubB1(1)
b2 = SubB2(2)
_ = b0.getOther() # Type checker => A, OK
_ = b1.getOther() # Type checker => A, Not OK
_ = b2.getOther() # Type checker => A, Not OK
I also tried reverting the order of B and _GenericB in SubB1 and SubB2, hoping to get the correct resolution of getOther, the MRO resolution fails: (TypeError: Cannot create a consistent method resolution order (MRO) for bases _GenericB, B'):
...
class SubB1(_GenericB[SubA1], B) # B after _GenericB raises MRO resolution TypeError
class SubB1(_GenericB[SubA2], B) # B after _GenericB raises MRO resolution TypeError
3. Separate base classes
The third attempt I made was to separate the generic and non generic part of the B class into separate base classes, having B, SubB1 and SubB2 inherit from them.
This however breaks inheritance relations:
...
class _NongenericB:
def __init__(self, b:int) -> None:
self.b:int = b
class _GenericB(Generic[Aclass]):
def getOther(self) -> Aclass:
return self._otherClass(a = self.b)
class B(_NongenericB, _GenericB[A]):
_otherClass = A
class SubB1(_NongenericB, _GenericB[SubA1]):
_otherClass = SubA1
class SubB2(_NongenericB, _GenericB[SubA2]):
_otherClass = SubA2
However:
_ = b0.getOther() # Type checker => A, OK
_ = b1.getOther() # Type checker => SubA1, OK
_ = b2.getOther() # Type checker => SubA2, OK
assert isinstance(b0, B) # OK
assert isinstance(b1, B) # Assertion error
assert isinstance(b2, B) # Assertion error
Also, I cannot add B to the base classes of SubB1 and SubB2 because I get in the same situation of point 2.
I am a bit at a loss here. What should I do?
One obvious solution would be to redefine the method in the subclasses and annotate the signatures manually, without using generics, but this is something that I would like to avoid, if possible, as there may be many methods working with this mechanism. Having the method defined in the base class once and for all would the ideal solution.
Any ideas?
Thank you!
Your attempt Making B a generic class was actually almost there. You only need the following modifications if you're using mypy or pyright (other type-checkers aren't up-to-spec and have non-standard workarounds):
typing_extensions.TypeVar if you're using Python < 3.13, then add a bound= and use type variable defaults:
Aclass = TypeVar('Aclass', bound='A', default='A')
bound=A allows correct signature resolution when you're calling self._otherClass(...) (see 2. below), because you'd want your type checker to use the class constructor of type(self)'s concrete type argument.default=A allows you to call B(0) without specifying the generic type of B - in this case the generic type of B in the call B(0) defaults to A._otherClass - you'll have to suppress a non-harmful error here:
class B(Generic[Aclass]):
_otherClass: type[T] = A # type: ignore[assignment]
... and you're done! See demonstrations below: