TL;DR

Use from typing import dataclass_transform

Motivation

其实动机很简单，众所周知对于 SQLAlchemy 1.4 想要 typing 可以安装 sqlalchemy2-stub，对于 < 1.4 也可以有 sqlalchemy-stub。然而对于最新的 SQLAlchemy >= 2.0，因为它自己有类型注释，但是又很少，所以还没有一个很好的解决方案。本文就是为了介绍一种我摸索出的解决方案，基本上可以完美解决 SQLAlchemy Model / ORM 的 typing 问题。

当然，如果你在手写 raw sql，那肯定是没办法自动弄好类型的，不要做梦了🚫。本文只针对使用了 SQLAlchemy ORM Model 的用户。

Implementation

基本思路是，Python 在 PEP 681 (>= Python 3.11) 当中为 typing 模块提供了一个 dataclass_transform decorator，可以将第三方的 class 标注为和原生的 dataclass 提供类似的功能：

Most type checkers, linters and language servers have full support for dataclasses. This proposal aims to generalize this functionality and provide a way for third-party libraries to indicate that certain decorator functions, classes, and metaclasses provide behaviors similar to dataclasses.

These behaviors include:

  • Synthesizing an __init__ method based on declared data fields. <-- good for us
  • Optionally synthesizing __eq__, __ne__, __lt__, __le__, __gt__ and __ge__ methods.
  • Supporting “frozen” classes, a way to enforce immutability during static type checking.
  • Supporting “field specifiers”, which describe attributes of individual fields that a static type checker must be aware of, such as whether a default value is provided for the field. <-- mostly good for us

总体来讲这个 decorator 提供了所有我们想要的功能，唯一的问题是 field: type = mapped_column(...) 会导致 type checker 认为 field 是 Optional 的。但总归基本的类型检查是能用的，甚至能兼容外键和 relationship，只是在初始化类的时候会有提供参数没写全的风险。如果感觉这样不够好，可以通过给每个 Model 编写一个 dummy __init__ 的方式来把 optional 干掉：

1
2
3
4
5
class A(Base):
    id: MappedColumn[str] = mapped_column(String, nullable=False, ...)

    def __init__(/, id: str):
        ...

实践上，我们为了方便，可以直接在 SQLAlchemy 的 Base class 上面使用这个 decorator，这样就不用每次定义一个 Model class 都要写一遍了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from sqlalchemy import Integer, String, Text, text, ForeignKey
from sqlalchemy.orm import DeclarativeBase, Mapped, MappedColumn, mapped_column, relationship

from typing import TYPE_CHECKING, dataclass_transform

if TYPE_CHECKING:
    # this gives any class that derives from Base some simple type hints
    @dataclass_transform()
    class Base(DeclarativeBase):
        pass
else:
    # dataclass_transform should never mess with runtime codes to avoid behavior changes
    class Base(DeclarativeBase):
        pass

class User(Base):
    id: MappedColumn[str] = mapped_column(String, nullable=False, primary_key=True)
    ...

class Example(Base):
    __tablename__ = "example"
    id: MappedColumn[int] = mapped_column(Integer, autoincrement=True, primary_key=True)
    user_id: MappedColumn[str] = mapped_column(String, ForeignKey("user.id"), nullable=False)

    user: Mapped[User] = relationship(User, backref="example")  # should not use Relationship[] here

example = Example(...)  # has autocompletion & basic type check
example.user.id  # has autocompletion & basic type check

这里 class Base(DeclarativeBase): 写法得到的 Base class 和 Base = declarative_base() 得到的 Base 在功能上是一样的，但是给了我们使用 dataclass_transform 的空间。

来源：https://blog.jiejiss.com/