How to add type annotations to SQLAlchemy models

For a long time, I’ve been advocating for using type annotations in languages like Python because of its benefits. To mention some, we can improve the robustness, testability, and scalability of our applications. Just as important, we can reduce the number of unit tests that we need to write. The goal is to reduce bugs at run-time. This technique can be used in languages with a strong type system.

Does Python have a strong type system? Of course! This type-safety example is a simple demonstration of this statement:

>>> foo = 5
>>> bar = 'hello'
>>> foo + bar
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for +: 'int' and 'str'
>>>

People tend to think that while Python is a dynamic "interpreted" language (spoiler alert: it's not), the type system is weak, with similar behavior to JavaScript's implicit type coercion:

Welcome to Node.js v16.13.2.
Type ".help" for more information.
> foo = 5
5
> bar = 'hello'
'hello'
> foo + bar
'5hello'
>

Going back to Python and its type annotations feature, you might have a medium-sized project that needs to use database storage. However, it's common to leave that part out of our safe and sound type annotations. But fear not, you can also annotate your models. In this example, I'm going to show you a small example using Flask and SQLAlchemy.

Some time ago, I built a tiny sample project to demonstrate TDD on Flask. Let’s use it as a starting point for this tutorial. This example will have a small model for storing notifications in a database and a couple of endpoints to get/post some data.

We’re going to use Flask-SQLAlchemy, which provides a quick integration between the framework and the ORM. So after setting up the initialization we can start writing a model:

from flask_sqlalchemy.model import DefaultMeta

from sqlalchemy import Integer, String, Datetime, Boolean

BaseModel: DefaultMeta = db.Model


class Notification(BaseModel):
    __tablename__ = "notifications"

    id = Column(Integer, primary_key=True, index=True)
    description = Column(String(100))
    email = Column(String(100), nullable=False)
    date = Column(DateTime, nullable=False)
    url = Column(String(200))
    read = Column(Boolean, default=False)

We can explain some details here:

1. Why the BaseModel definition? Answer: We're using mypy for the static type checking. Defining a model as Flask-SQLAlchemy suggests causes an issue on that line because the declarative base db.Model is defined dynamically (at runtime) contrary to what a static checker expects.

2. Flask-SQLAlchemy suggests that we can use db.Column or db.Integer to define our properties, but we won't follow that convention. The reason behind this decision is because Flask-SQLAlchemy does a lot of work dynamically at runtime that can confuse the static type checking system. To avoid that, we are using SQLAlchemy classes directly.

The static type checker combination will be mypy and sqlalchemy-stubs plugin.

Just install them with pip:

pip install mypy
pip install sqlalchemy-stubs

Next, create a mypy.ini file in the root of the project:

[mypy]
python_version = 3.9
ignore_missing_imports = True
plugins = sqlmypy

Then you're ready to do some static type check analysis to the Python file:

mypy app.py

We're going to expand this model a bit. First, we need a method to "serialize" the model into a dict so it can be returned in the endpoint. Also, we need a couple of methods to get all the records filtered by some criteria.

We will define a NotificationModel dictionary definition.

class NotificationModel(TypedDict):
    id: int
    description: Optional[str]
    email: str
    date: datetime
    url: Optional[str]
    read: bool

Then, our basic serialization method for the Notification model will look like this.

def to_dict(self) -> NotificationModel:
    return {
        "id": self.id,
        "email": self.email,
        "description": self.description,
        "date": self.date,
        "url": self.url,
        "read": self.read,
    }

Here you can see the type annotations in action. Either with a properly configured editor (like VSCode) or with the mypy command, types should match between what to_dict returns and the NotificationModel definition, otherwise the analysis will fail.

Example: if the to_dict method returns a string value in the id property

def to_dict(self) -> NotificationsModel:
        return {
            "id": str(self.id),  # <-- should not be str

then the error we get will be the following.

app.py:38: error: Incompatible types (expression has type "Optional[str]", TypedDict item "id" has type "int")
Found 1 error in 1 file (checked 1 source file)

Next, we will implement a couple of class methods that will return all the records in the table and another to get only the unread notifications:

@classmethod
def get_all(cls) -> list[Notification]:
    return db.session.query(cls).all()

@classmethod
def get_unread(cls) -> list[Notification]:
    return db.session.query(cls).filter(Notification.read.is_(False)).all()

A key point here in order for a method to return its same class is to have this line on top of the module. See postponed evaluation of annotations for more details. py from __future__ import annotations

At this point, we have our model ready to be used in a presentation layer.

Let's implement a /notifications endpoint that accepts two methods, GET to list all notifications, and POST to save a new one with form data:

@app.route("/notifications", methods=["GET", "POST"])
def notifications() -> Response:
    if request.method == "POST":
        new_notification = Notification(
            **dict(request.form, date=datetime.fromisoformat(request.form["date"]))
        )
        db.session.add(new_notification)
        db.session.commit()

    notifications = Notification.get_all()
    return jsonify([notification.to_dict() for notification in notifications])

We implemented another route to get all unread notifications but it's pretty similar to the GET section of the above endpoint. You can review the entire sample project in this repository:

https://github.com/po5i/flask-mini-orm

The usage of these endpoints is beyond the scope of this post. They were only implemented to demonstrate and test the SQLAlchemy Notification model. I'd also like to point out that this example has everything implemented in a single module. We recommend you to (1) split your models and (2) use a migration tool such as Alembic. Last but not least, (3) test your logic with unit tests. We didn’t add them but feel free to propose some as pull requests. In the meantime, we configured the CI workflow with the mypy type checking.

Are you using Django? There are also some stubs ready to annotate your models: django-stubs and djangorestframework-stubs can suit your needs. Let us know if you want to read some content about them.

Thanks for reading :bow: