Python 3 and Generic Classes for IDE type-hinting

With Python 3 type annotations comes the ability to declare type Generics, which is a kind of way of making a template out of a piece of code for a specific type. For example, you can declare a Stack class that can implement a stack interface, but restrict it to only containing values of a certain type (at least with type annotations; nothing at runtime prevents multiple types from being passed). Rather than hard-coding the individual types that are allowed, a “generic” type can be declared; a variable type, if you will.

Creating a Generic variable

The typing module that comes built-in to Python 3 provides a TypeVar class that will create a Generic type variable. It is created like so:

from typing import TypeVar

T = TypeVar('T')

Using the Generic variable within a class

In order to use this generic variable within a class, the Generic class must be imported and inherited within the base class members of the class. Then, any methods can use the generic variable, T, as input values or output values:

from typing import Generic

class MyClass(Generic[T]):
    def do_thing(self, items: [T]) -> Optional[T]:
        for item in items:
            value = item.some_method_on_t()
            if some_condition(value):
                return value
        return None

Creating a generic base class with a concrete subclass

For my projects at work I’m using SQLAlchemy (SQLA) to interact with a database and I’ve created a base QueryDao class that wraps around SQLA’s Query class and provides business logic-specific query support while maintaining the query builder pattern of SQLA (each method returns self). The QueryDao class is a generic class that should work for any database model. As such, I have declared it like so:

from typing import Generic TypeVar, Type
from models import Base  # SQLA's declarative_base class

ModelType = TypeVar("ModelType", bound=Base)


class QueryDao(Generic[ModelType]):
    Model: Type[ModelType] = None

    def __init__(self, session):
        if self.Model is None:
            raise MisconfiguredDaoException(
                self, "model", "This should be set to a Model class.",
            )

        self._session = session
        self.query = self._new_query()
        self.offset = None
        self.size = None

    def _new_query(self):
        return self._session.query(self.Model)

    def build(self):
        """
        Finalizes the query and returns it. 
        This resets the Dao to a blank state with 
        a new query.
        """
        query = self.query

        if self.size:
            query = query.limit(self.size)

        if self.offset:
            query = query.offset(self.offset)

        self.query = self._new_query()
        self.reset()

        return query

    def get(self, entity_id, should_error=True) -> ModelType:
        """
        Gets the Model by its primary key.

        This uses the session cache to prevent a DB fetch.

        If should_error is set, this will raise an exception
        if not found (default).
        """
        result = self._new_query().get(entity_id)
        if not result and should_error:
            raise NoResultFound("No result was found.")

        return result

    def all(self) -> [ModelType]:
        """
        Finalizes the query and returns all rows in the 
        database by this query.
        """
        return self.build().all()

The specific things to note here are:

  • The QueryDao class inherits from Generic[ModelType]
  • The Model class-level attribute has a type of Type[ModelType], indicating it is a class of type ModelType
  • The .get() and .all() methods return ModelType and [ModelType] respectively.

Now, I can subclass and supply a concrete SQLA Model class so that all of the type hints for Model, .get(), and .all() are supplied with the correct value within editors like PyCharm:

class MyModelDao(QueryDao[MyModel]):
    # other specific methods here


# then, in another file where I can use the MyModelDao
dao = MyModelDao(session)
dao.get(1234)  # type => MyModel

Because QueryDao inherits from Generic[ModelType], we can use indexing to provide a concrete model type for the subclass. This is done with QueryDao[MyModel]syntax. Now, any method that we use on MyModelDao that is inherited from QueryDao will show the correct return type, e.g. .get() will show a return type of MyModel which then allows us to use attributes and methods on the value with proper type-hinting!

2 thoughts on “Python 3 and Generic Classes for IDE type-hinting”

  1. This is great, I was trying to build something similar but got stuck. Specifically, I couldn’t figure out line 8:
    Model: Type[ModelType] = None
    but could only make it work with a class that takes a generic type AND a constructor argument.

    Thank you!

    Reply
  2. Awesome, thank you for the great example. I was trying to do the same but couldn’t work out how to get the generic type into an instance attribute. So I had to pass the model type as the generic type as well as a constructor argument, which was ugly. This is way better!

    Reply

Leave a Reply