Python SQLAlchemy 2.0 使用 dataclass_transform 非必需字段类型

我刚刚在一个新项目上安装了 SQLAlchemy 2.0，我正在尝试使我的模型尽可能类型安全。

通过使用 @typing_extensions.dataclass_transform ，我已经能够实现我想要实现的大部分目标类型检查，但是当前所有字段都被标记为不需要。

例如：

@typing_extensions.dataclass_transform(kw_only_default=True)
class Base(DeclarativeBase):
    pass

class TestModel(Base):
    __tablename__ = "test_table"
    name: Mapped[str]
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    external_id: Mapped[int] = mapped_column(
        ForeignKey("external.id"), nullable=False
    )

def test_test_model(session: Session) -> None:
    TEST_NAME = "name"
    external = External()
    session.add(external)
    session.commit()
    model1 = TestModel() # Intellisense shows error because "name" is required

    model2 = TestModel(name=TEST_NAME, external_id=external.id). # no error
    session.add(model2)
    session.commit() # model commits successfully

    model3 = TestModel(name=TEST_NAME) # No intellisense error, despite "external_id" being required
    session.add(model3)
    session.commit(). # error when saving because of missing "external_id"

在上面的示例中，如何将 external_id 的类型设置为必需？

问题在于，当在 mapped_column() 中设置 nullable=False 时，这只会影响数据库级别的约束，而不会影响 Python 类型提示。

为了使 external_id 在类型提示中成为必需项，需要在 dataclass 定义中提供默认值，即使该默认值在数据库级别无关紧要。以下是一些方法：

1. 使用 Field 和 default_factory :

from sqlalchemy import ForeignKey
from sqlalchemy.orm import Mapped, mapped_column, DeclarativeBase
from typing import Optional
from sqlalchemy.orm import Session
from typing_extensions import dataclass_transform

@dataclass_transform(kw_only_default=True)
class Base(DeclarativeBase):
    pass

class External(Base):
    __tablename__ = "external"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)

class TestModel(Base):
    __tablename__ = "test_table"
    name: Mapped[str] 
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    external_id: Mapped[int] = mapped_column(ForeignKey("external.id"), nullable=False)

def test_test_model(session: Session) -> None:
    TEST_NAME = "name"
    external = External()
    session.add(external)
    session.commit()

    # 使用默认工厂函数提供默认值
    model1 = TestModel(name=TEST_NAME, external_id=external.id) 
    session.add(model1)
    session.commit() 

2. 使用 __post_init__ :

from sqlalchemy import ForeignKey
from sqlalchemy.orm import Mapped, mapped_column, DeclarativeBase
from typing import Optional
from sqlalchemy.orm import Session
from typing_extensions import dataclass_transform

@dataclass_transform(kw_only_default=True)
class Base(DeclarativeBase):
    pass

class External(Base):
    __tablename__ = "external"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)

class TestModel(Base):
    __tablename__ = "test_table"
    name: Mapped[str]
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    external_id: Mapped[int] = mapped_column(ForeignKey("external.id"), nullable=False)

    def __post_init__(self):
        if self.external_id is None:
            raise ValueError("external_id is required")

def test_test_model(session: Session) -> None:
    TEST_NAME = "name"
    external = External()
    session.add(external)
    session.commit()

    model1 = TestModel(name=TEST_NAME, external_id=external.id)
    session.add(model1)
    session.commit() 

这两种方法都可以在类型提示中将 external_id 设置为必需项，同时仍然允许在数据库级别强制执行非空约束。选择哪种方法取决于的偏好和具体用例。