How to add type annotations to SQLAlchemy models
Tags: python, orm, flask, typechecking
Published: 2022-03-21 13:07:02

For a long time, I’ve been advocating for using [type annotations](https://www.stackbuilders.com/blog/using-types-in-python-with-mypy/) 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 of 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](https://flask.palletsprojects.com/) and [SQLAlchemy](https://www.sqlalchemy.org/). 

Some time ago, I built a tiny [sample project to demonstrate TDD on Flask](https://github.com/po5i/flask-mini-tests). 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](https://flask-sqlalchemy.palletsprojects.com/en/2.x/), which provides a quick integration between the framework and the ORM. So after setting up the initialization we can start writing a model:

```py
from flask_sqlalchemy.model import DefaultMeta


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](https://mypy.readthedocs.io/en/stable/) 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 (on 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 on 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](https://github.com/dropbox/sqlalchemy-stubs).

Just install them with pip:

```bash
pip install mypy
pip install sqlalchemy-stubs
```

And create a `mypy.ini` file on the root of the project:

```ini
[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:

```bash
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.

```py
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.

```py
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 on the `id` property

```py
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 on the table and another to get only the unread notifications:

```py
@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](https://docs.python.org/3/whatsnew/3.7.html#pep-563-postponed-evaluation-of-annotations) for more details.
> ```py
> from __future__ import annotations
> ```

At this point, we have our model ready to be used on 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:

```py
@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:

{% github 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 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](https://alembic.sqlalchemy.org/en/latest/). 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.

I encourage you to use a better tool to implement REST APIs like I did on a previous post:

{% embed https://dev.to/po5i/how-to-set-up-a-rest-api-in-flask-in-5-steps-5b7d %}

Are you using Django? There are also some stubs ready to annotate your models: [django-stubs](https://github.com/typeddjango/django-stubs) and [djangorestframework-stubs](https://github.com/mkurnikov/djangorestframework-stubs) than can suit your needs. Let me know if you want to read some content about them.

Thanks for reading :bow:

Edit: Article also cross-posted [here](https://www.stackbuilders.com/blog/how-to-add-type-annotations-to-sqlalchemy-models/).