如何从SQLAlchemy声明式类自动生成TypedDict类型定义？

我定义了若干代表数据库表的ORM类（基于SQLAlchemy 2和Python 3.9），示例代码如下：

import datetime
from typing import Optional
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

class Base(DeclarativeBase):
    pass

# ORM class
class Predictions(Base):
    __tablename__ = 'predictions'
    uid: Mapped[str] = mapped_column(primary_key=True)
    event_date: Mapped[datetime.date]
    location: Mapped[Optional[str]]

# ... other tables

由于我会将查询结果以字典形式返回，希望为每个表生成对应的TypedDict类型定义，示例如下：

# desired TypedDict
class PredictionsDict(TypedDict):
    uid: str
    event_date: datetime.date
    location: Optional[str]

请问如何自动生成可被智能感知识别的TypedDict类型定义，避免手动维护多个类的同步？

方法1：利用SQLAlchemy inspect工具动态生成 #

通过SQLAlchemy的inspect函数提取ORM类的字段元数据，再自动生成对应TypedDict：

import datetime
from typing import Optional, TypedDict, get_args
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, inspect

class Base(DeclarativeBase):
    pass

class Predictions(Base):
    __tablename__ = 'predictions'
    uid: Mapped[str] = mapped_column(primary_key=True)
    event_date: Mapped[datetime.date]
    location: Mapped[Optional[str]]

def generate_typed_dict(orm_cls) -> type[TypedDict]:
    mapper = inspect(orm_cls)
    fields = {}
    # 遍历ORM类的所有列字段
    for column in mapper.columns:
        type_hint = orm_cls.__annotations__[column.name]
        # 从Mapped类型中提取原始Python类型
        if hasattr(type_hint, '__origin__') and type_hint.__origin__ is Mapped:
            fields[column.name] = get_args(type_hint)[0]
        else:
            fields[column.name] = type_hint
    # 动态创建TypedDict类
    return type(f"{orm_cls.__name__}Dict", (TypedDict,), fields)

# 生成并使用TypedDict
PredictionsDict = generate_typed_dict(Predictions)

# IDE可智能感知字段类型
sample_dict: PredictionsDict = {
    "uid": "123",
    "event_date": datetime.date.today(),
    "location": "NYC"
}

方法2：通过元类自动绑定TypedDict #

给Base类添加自定义元类，在ORM类定义时自动生成并绑定对应的TypedDict：

import datetime
from typing import Optional, TypedDict, get_args
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

class BaseMeta(type(DeclarativeBase)):
    def __new__(cls, name, bases, namespace):
        orm_cls = super().__new__(cls, name, bases, namespace)
        # 跳过Base类本身
        if name == "Base":
            return orm_cls
        # 提取Mapped类型对应的原始类型
        fields = {}
        for attr_name, type_hint in namespace.get('__annotations__', {}).items():
            if hasattr(type_hint, '__origin__') and type_hint.__origin__ is Mapped:
                fields[attr_name] = get_args(type_hint)[0]
        # 创建TypedDict并绑定到ORM类
        typed_dict_cls = type(f"{name}Dict", (TypedDict,), fields)
        setattr(orm_cls, f"{name}Dict", typed_dict_cls)
        return orm_cls

class Base(DeclarativeBase, metaclass=BaseMeta):
    pass

class Predictions(Base):
    __tablename__ = 'predictions'
    uid: Mapped[str] = mapped_column(primary_key=True)
    event_date: Mapped[datetime.date]
    location: Mapped[Optional[str]]

# 直接使用ORM类上绑定的TypedDict
sample_dict: Predictions.PredictionsDict = {
    "uid": "456",
    "event_date": datetime.date(2024, 1, 1),
    "location": None
}

注意事项 #

• 两种方法均支持Python 3.9+和SQLAlchemy 2.x，生成的TypedDict可被主流IDE（PyCharm、VS Code）智能感知
• 若ORM类包含relationship等非列字段，需在遍历字段时过滤掉，避免生成无效的TypedDict键
• 对于SQLAlchemy自定义类型（如ARRAY、JSON），需确保get_args能正确提取对应的Python类型

内容的提问来源于stack exchange，提问作者jonah0