CODEWITHSUNDEEP
pythonpython-3.xpython-typingmypy
Trevor Gross
Trevor Gross

Reputation: 678

Type Hinting: Use type of a class member as function return type (for inheritance)

What is the correct way to reuse the type of a class member to type hint other items in the class? As an example:

from typing import Type

class Model:
    pass

class ChildModel:
    childvar = "Child Model"

class Base:
    var: Type[Model]

    def fn(self) -> ??:
        return self.var

class Child(Base):
    var = ChildModel

    def new_fn(self):
        x = self.fn()  # Type of x should be "ChildModel"
        print(x.childvar)

Child().new_fn() # Prints "Child Model" successfully

I am looking for what would work to replace ?? such that the return type of fn() can be inferred for all child classes.

MyPy does not accept changing ?? to Type[Model] to match Base.var: Incompatible types in assignment (expression has type "Type[ChildModel]", base class "Base" defined the type as "Type[Model]" (though it is possible I made a mistake here). Even if this were allowed, this would allow Base.fn() to return any Model or Model subclass, not strictly the type of var (as defined in a child of Base)

Something like T = TypeVar("T", bound=Type[Model]) seems disallowed without generics, which don't seem quite applicable since the type can be inferred without generic-style specification. I think the solution would likely also work to type hint method arguments, method-local variables, and other class member variables.

What is the best way to do this (if possible)?

Edit: adding clarification, corrected issue with code

Upvotes: 5

Views: 3151

Answers (2)

Farhaan S.
Farhaan S.

Reputation: 104

Though this probably fails the "Explicit is better than Implicit" test, I suppose this will get you what you want while avoiding typing in two places. In this case, rather than defining var on the Child, the var is pulled from the annotation.

Tested on Python 3.10

import typing
from typing import Generic, TypeVar

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


class Model:
    pass


class ChildModel(Model):
    childvar = "Child Model"


class Base(Generic[T]):
    @classmethod
    @property
    def var(cls) -> type[T]:
        for superclass in cls.__orig_bases__:
            if getattr(superclass, "__origin__", None) == Base:
                return typing.get_args(superclass)[0]

    def fn(self) -> type[T]:
        return self.var


class Child(Base[ChildModel]):
    def new_fn(self):
        x = self.fn()  # Type of x is type["ChildModel"]
        print(x.childvar)

Upvotes: 1

Paweł Rubin
Paweł Rubin

Reputation: 3350

This can be accomplished with Generics.

from typing import Generic, TypeVar

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


class Model:
    pass


class ChildModel(Model):
    childvar = "Child Model"


class Base(Generic[T]):
    var: type[T]

    def fn(self) -> type[T]:
        return self.var


class Child(Base[ChildModel]):
    var = ChildModel

    def new_fn(self):
        x = self.fn()  # Type of x is type["ChildModel"]
        print(x.childvar)


Child().new_fn()

Upvotes: 3

Related Questions