How to annotate parent classmethod so that child methods return an instance of themselves?

Question:

I have a parent class BaseBlob that has several child classes LimitedProofId, ProofId, TxId. The parent class implement a deserialize class method that should return an instance of itself.

I also have a Delegation class that takes a LimitedProofId. I specifically want mypy to error if I pass by mistake another child instance of BaseBlob such as a ProofId or a TxId.

from __future__ import annotations
from io import BytesIO


class BaseBlob:
    def __init__(self, data: bytes):
        self.data = data

    @classmethod
    def deserialize(cls, stream: BytesIO) -> BaseBlob:
        return cls(stream.read(32))


class LimitedProofId(BaseBlob):
    pass


class TxId(BaseBlob):
    pass


class Delegation:
    def __init__(self, ltd_id: LimitedProofId):
        self.ltd_id = ltd_id

    def deserialize(self, stream: BytesIO) -> Delegation:
        ltd_id = LimitedProofId.deserialize(stream)
        return Delegation(ltd_id)

mypy shows an error for this code, because if thinks LimitedProofId.deserialize returns a BaseBlob.

error: Argument 1 to "Delegation" has incompatible type "BaseBlob"; expected "LimitedProofId"  [arg-type]

I have seen answer to similar questions that use a T = TypeVar('T', bound='BaseBlob') to achieve a type annotation that allows child classes, but if I do that I need to specify T both for the return type of BaseBlob.deserialize and the first parameter of Delegation.__init__, which defeat my purpose of type safety for the latter.

Is there a way to achieve what I want to do, without having to reimplement deserialize on all child classes?

Asked By: PiRK

||

Answers:

Python 3.11 introduces the Self type hint for this. (PEP 673 describes in more detail the code Self is intended to simplify, if you haven’t upgraded to 3.11 yet.)

from typing import Self


class BaseBlob:
    def __init__(self, data: bytes):
        self.data = data

    @classmethod
    def deserialize(cls, stream: BytesIO) -> Self:
        return cls(stream.read(32))
Answered By: chepner

You want to express that deserialize returns an instance of the class that it is bound to.

Python >=3.9, <3.11

...
from typing import TypeVar

T = TypeVar("T", bound="BaseBlob")

class BaseBlob:
    ...

    @classmethod
    def deserialize(cls: type[T], stream: BytesIO) -> T:
        return cls(stream.read(32))

These changes make the code you posted above perfectly type safe and mypy --strict agrees.

Calling from a child class like LimitedProofId.deserialize binds the method, so that cls will be LimitedProofId, which from a typing perspective in turn binds T in type[T] accordingly.


Python <3.9

Like the above, but import Type from typing and replace the type[T] annotation with Type[T].


Python >=3.11

What @chepner said. Should work soon with mypy.


Clarification

I don’t understand, what you mean: (highlighted)

I have seen answer to similar questions that use a T = TypeVar('T', bound='BaseBlob') to achieve a type annotation that allows child classes, but if I do that I need to specify T both for the return type of BaseBlob.deserialize and the first parameter of Delegation.__init__, which defeat my purpose of type safety for the latter.

What does __init__ have to do with this?

If I failed to understand your intention, please elaborate and I will do my best to adjust the answer.


Addendum

Since you seem to be unsure of how these annotations affect how type checkers infer the relevant types, here is a full demonstration:

from __future__ import annotations
from io import BytesIO
from typing import TypeVar

T = TypeVar("T", bound="BaseBlob")

class BaseBlob:
    def __init__(self, data: bytes):
        self.data = data

    @classmethod
    def deserialize(cls: type[T], stream: BytesIO) -> T:
        return cls(stream.read(32))

class LimitedProofId(BaseBlob):
    pass

class TxId(BaseBlob):
    pass

class Delegation:
    def __init__(self, ltd_id: LimitedProofId):
        self.ltd_id = ltd_id

    @classmethod
    def deserialize(cls, stream: BytesIO) -> Delegation:
        ltd_id = LimitedProofId.deserialize(stream)
        return cls(ltd_id)

if __name__ == "__main__":
    ltd = LimitedProofId(b"xyz")
    d = Delegation.deserialize(BytesIO(b"abc"))
    reveal_type(d.ltd_id)
    reveal_type(LimitedProofId.deserialize(BytesIO()))
    reveal_type(TxId.deserialize(BytesIO()))

The mypy --strict output is without errors:

note: Revealed type is "LimitedProofId"
note: Revealed type is "LimitedProofId"
note: Revealed type is "TxId"
Success: no issues found in 1 source file
Answered By: Daniil Fajnberg