Python type annotation improvement with Generic and TypeVar

Posted at — Oct 15, 2021

In this article, I would like to share my favorite way of type annotating classes using typing modules with Generic and TypeVar.

Checker installation

First, to check python typing you can use mypy library or any IDE/editor with type checker like Pycharm.

pip install mypy

Generic

It is an abstract base class for generic types. A generic type is typically declared by inheriting from an instantiation of this class with one or more type variables

TypeVar

It is a type variable. Type variables exist primarily for the benefit of static type checkers. They serve as the parameters for generic types as well as for generic function definitions.

Showcase

As an example, we will build a base class for our database repository, and then we will use it.

First, let’s define our base repository:

import abc

from schemas.base import BaseSchema
from tables.base import BaseTable


class BaseRepository(metaclass=abc.ABCMeta):
    @property
    @abc.abstractmethod
    def _schema(self) -> Type[BaseSchema]:
        ...
    
    @property
    @abc.abstractmethod
    def _table(self) -> Type[BaseTable]:
        ...
    
    def create(self, schema: BaseSchema) -> Type[BaseTable]:
        table = self._table(**schema.__dict__)
        return table

Now, let’s build the test repository:

from schemas.base import BaseSchema
from tables.base import BaseTable


@dataclass()
class UserSchema(BaseSchema):
    id: int
    name: str


@dataclass()
class UserTable(BaseTable):
    id: int
    name: str


class UserRepository(BaseRepository):
    @property
    def _table(self) -> UserTable:
        return UserTable

    @property
    def _schema(self) -> UserSchema:
        return UserSchema

    def perform_create(self) -> UserTable:
        schema = UserSchema(id=1, name="Piotr")
        return self.create(schema)

The problem with the above code is that every time we want to run the create method it will return the BaseTable type. You will see the error:

Expected type 'UserTable', got 'BaseTable' instead

We can fix this by introducing TypeVar with Generic:

import abc 
from typing import TypeVar, Generic

from schemas.base import BaseSchema
from tables.base import BaseTable


SCHEMA = TypeVar("SCHEMA", bound=BaseSchema)
TABLE = TypeVar("TABLE", bound=BaseTable)


class BaseRepository(Generic[SCHEMA, TABLE], metaclass=abc.ABCMeta):
    @property
    @abc.abstractmethod
    def _schema(self) -> Type[SCHEMA]:
        ...
    
    @property
    @abc.abstractmethod
    def _table(self) -> Type[TABLE]:
        ...
    
    def create(self, schema: SCHEMA) -> TABLE:
        table = self._table(**schema.__dict__)
        return table

New implementation of the UserRepository should inherit now from the base with the definition of the types we want to use:

from schemas.base import BaseSchema
from tables.base import BaseTable


@dataclass()
class UserSchema(BaseSchema):
    id: int
    name: str


@dataclass()
class UserTable(BaseTable):
    id: int
    name: str


class UserRepository(BaseRepository[UserSchema, UserTable]):
    @property
    def _table(self) -> UserTable:
        return UserTable

    @property
    def _schema(self) -> UserSchema:
        return UserSchema

    def perform_create(self) -> UserTable:
        schema = UserSchema(id=1, name="Piotr")
        return self.create(schema)

Introducing the above change we are eliminating the error and making code cleaner and better type hinted.

Summary

This is one of many cases where this approach can be used. Feel free to experiment and play on your own. I have already used it in my previous tutorial on the example with SQLAlchemy, FastAPI, and Pydantic, check it on my github.

rogulski.it

My name is Piotr, a passionate pythonista and this is my blog!