How to add type hints to your SQLAlchemy models?

We wrote about adding type hints to the SQLAlchemy models some time ago and fortunately SQLAlchemy and Flask-SQLAlchemy's new version simplify the process of getting typed models and the connection with a database. In this blog post, we’ll revamp the blog about adding types but highlighting how SQLAlchemy allows us to have types models effortlessly.

NOTE: In this post, we’ll use pyright for the static analysis because it offers some type inference. Type inference will save us some time since we won’t have to define explicit types on every function. Flask-SQLAlchemy

Flask-SQLAlchemy

We will use Flask-SQLAlchemy version 3.1.x and SQLAlchemy version 2.x. The new versions of these libraries have features that will help us have typed models as in the original post but with a different approach that limits the purpose of each library. Flask-SQLAlchemy will continue as a wrapper to automatically set up the engine and session to manage data from the database in a Flask Application, and SQLAlchemy will be the tool that helps us define models through improved support for Python type hints and data classes. The start point is going to be the model definition, and for this part, we need an instance of SQLAlchemy and a subclass of DeclarativeBase that are defined in Flask-SQLAlchemy and SQLAlchemy, respectively:

from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.orm import DeclarativeBase

class Base(DeclarativeBase):
  pass

db = SQLAlchemy(model_class=Base)

The model will start as a subclass of db.Model and use the new Mapped class and mapped_colum helper from SQLAlchemy to define the attributes that will map the columns of our tables. The following example shows them in action:

from datatime import datime
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

class Notification(db.Model):
    __tablename__ = "notifications"

    id: Mapped[int] = mapped_column(primary_key=True, index=True)
    description: Mapped[str]
    email: Mapped[str] = mapped_column(nullable=False)
    date: Mapped[datetime] = mapped_column(nullable=False)
    url: Mapped[str]
    read: Mapped[bool] = mapped_column(default=False)

We can highlight the following from the model definition:

• We’re defining the model attributes with the Mapped generic class and Python types to define the model.
• The mapped_column helper specifies primary keys, indexes, nullables, and default values. This helper can also be used to set other settings like uniqueness, autoincrement, etc.

Once the model is defined, db.create_all() will create everything (database, tables, columns, indices, etc) for us within the app.app_context() context:

with app.app_context():
       db.create_all()

Let’s connect the Application with our database, which is quite straightforward and self-explanatory:

from flask import Flask

app = Flask(__name__)

app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///project.db"

db.init_app(app)

Next, the query implementation uses db.session.query as a common SQLAlchemy query. But, here comes one of the main benefits of using pyright: We won’t need to write down the type annotations of the functions, they will be inferred!

def get_all():
    return db.session.query(Notification).all()

def get_unread():
    return db.session.query(Notification).filter(Notification.read.is_(False)).all()

pyright users won’t have to write the function signature because the static checker will infer the type. Don’t believe me? Let’s check it out with the reveal_type helper from the typing package:

pyright app.py

reveal_type(get_all)
...: Type of "get_all" is "() -> List[Notification]"

reveal_type(get_unread)
...: Type of "get_unread" is "() -> List[Notification]"

Now, let’s put it all together in an application. We’ll use the NotificationModel TypedDict to represent the response of our endpoint in a primitive type:

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

The endpoint to add notifications will have the following structure:

@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 = get_all()
    return jsonify([notification.to_dict() for notification in notifications])

The full implementation is in the following repository, try it out! Type hints reduce feedback loops and move faster without the fear of breaking things. Static checkers will alert us in cases that we’re doing something that doesn’t make sense at the type level. For example:

get_all() - get_unread()

flask-sqlalchemy-typing/app/main.py:N:1 - error: Operator "-" not supported for types "List[Notification]" and "List[Notification]" (reportOperatorIssue)

Conclusion

The combination of Flask-SQLAlchemy and SQLAlchemy's new features, paired with a static type checker like Pyright, revolutionizes how we manage types in our models. By eliminating the need for additional dependencies like sqlalchemy-stubs and reducing configuration overhead, developers can now enjoy a streamlined experience. The instant feedback on type correctness empowers developers to catch errors early, saving time on debugging and manual testing. This improved workflow not only enhances productivity but also makes coding more efficient and enjoyable.