将ORM对象扁平化并映射为Pydantic模型的适配方案（支持FastAPI响应场景）

最近我在做项目时遇到一个需求：用Pydantic配合FastAPI把SQLAlchemy的ORM数据转换成JSON输出，需要对ORM模型做扁平化和字段重映射，去掉JSON里不必要的嵌套层级。举个直观的例子：

原始的ORM转换输出是这样的：

{"id": 1, "billing": [ {"id": 1, "order_id": 1, "first_name": "foo"}, {"id": 2, "order_id": 1, "first_name": "bar"} ] }

而我想要的简洁输出是：

{"id": 1, "name": ["foo", "bar"]}

处理纯字典输入的可行方案 #

一开始处理纯字典数据的时候，我通过重写Pydantic模型的__init__方法轻松实现了字段映射，代码如下：

from pydantic import BaseModel

# 测试用字典输入
order_dict = {"id": 1, "billing": {"first_name": "foo"}}
# 期望输出: {"id": 1, "name": "foo"}

class Order_Model_For_Dict(BaseModel):
    id: int
    name: str = None
    class Config:
        orm_mode = True

    def __init__(self, **kwargs):
        print( "kwargs for dictionary:", kwargs )
        # 这里能拿到完整的字典参数，直接做映射
        kwargs["name"] = kwargs["billing"]["first_name"]
        super().__init__(**kwargs)

print(Order_Model_For_Dict.parse_obj(order_dict)) # 输出：id=1 name='foo'

这段代码直接就能运行，完全符合预期。

遇到的坑：ORM对象转换时__init__不生效 #

但当我换成SQLAlchemy的ORM对象时，这个方法直接失效了——模型的__init__方法根本不会被调用。比如下面这个测试代码，最终输出的name字段还是None，完全没达到预期：

from pydantic import BaseModel, root_validator
from typing import List
from sqlalchemy.orm import relationship
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from pydantic.utils import GetterDict

Base = declarative_base()

# 定义ORM模型
class BillingOrm(Base):
    __tablename__ = "billing"
    id = Column(Integer, primary_key=True, nullable=False)
    order_id = Column(ForeignKey("orders.id", ondelete="CASCADE"), nullable=False)
    first_name = Column(String(20))

class OrderOrm(Base):
    __tablename__ = "orders"
    id = Column(Integer, primary_key=True, nullable=False)
    billing = relationship("BillingOrm")

# Pydantic模型
class Billing(BaseModel):
    id: int
    order_id: int
    first_name: str
    class Config:
        orm_mode = True

class Order(BaseModel):
    id: int
    name: List[str] = None
    # billing: List[Billing] # 取消注释可验证关联关系正常
    class Config:
        orm_mode = True

    def __init__(self, **kwargs):
        # 使用from_orm解析ORM对象时，这段代码根本不会执行
        print("kwargs for orm:", kwargs)
        kwargs["name"] = kwargs["billing"]["first_name"]
        super().__init__(**kwargs)

# 构造测试用ORM对象
billing_orm_1 = BillingOrm(id=1, order_id=1, first_name="foo")
billing_orm_2 = BillingOrm(id=2, order_id=1, first_name="bar")
order_orm = OrderOrm(id=1)
order_orm.billing.append(billing_orm_1)
order_orm.billing.append(billing_orm_2)

# 转换为Pydantic模型
order_model = Order.from_orm(order_orm)
print(order_model) # 输出：id=1 name=None，完全不是期望的['foo','bar']

而且这个场景和FastAPI里用response_model的写法完全一致，比如我们的接口通常是这样写的：

@router.get("/orders", response_model=List[schemas.Order])
async def list_orders(db: Session = Depends(get_db)):
    return get_orders(db)

所以必须找到适配这种场景的解决方案。

尝试root_validator遇到的问题 #

后来我试着用root_validator来处理，虽然能提取到想要的名称列表，但没法给values赋值——因为当输入是ORM对象时，values是GetterDict类型，不支持item赋值：

@root_validator(pre=True)
def flatten(cls, values):
    if isinstance(values, GetterDict):
        names = [ billing_entry.first_name for billing_entry in values.get("billing") ]
        print(names) # 这里能正确打印出['foo', 'bar']
        # values["name"] = names # 直接报错：'GetterDict' object does not support item assignment
    return values

可行的解决方案 #

经过一番折腾，我找到了两种靠谱的解决方法，都能完美适配FastAPI的响应模型场景：

方法一：在root_validator中转换GetterDict为普通字典 #

这种方法比较直接，在验证器里把GetterDict转成普通字典后再添加字段：

from pydantic import BaseModel, root_validator
from typing import List
from sqlalchemy.orm import relationship
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from pydantic.utils import GetterDict

Base = declarative_base()

# ORM模型定义
class BillingOrm(Base):
    __tablename__ = "billing"
    id = Column(Integer, primary_key=True, nullable=False)
    order_id = Column(ForeignKey("orders.id", ondelete="CASCADE"), nullable=False)
    first_name = Column(String(20))

class OrderOrm(Base):
    __tablename__ = "orders"
    id = Column(Integer, primary_key=True, nullable=False)
    billing = relationship("BillingOrm")

# 修正后的Pydantic模型
class Order(BaseModel):
    id: int
    name: List[str] = None
    class Config:
        orm_mode = True

    @root_validator(pre=True)
    def flatten_billing(cls, values):
        # 处理ORM对象转换时的GetterDict输入
        if isinstance(values, GetterDict):
            # 转成普通字典才能修改
            value_dict = dict(values)
            # 提取billing列表中的first_name字段
            value_dict["name"] = [b.first_name for b in value_dict.get("billing", [])]
            return value_dict
        # 兼容普通字典输入的情况
        elif "billing" in values:
            values["name"] = [b["first_name"] for b in values.get("billing", [])]
        return values

# 测试
billing_orm_1 = BillingOrm(id=1, order_id=1, first_name="foo")
billing_orm_2 = BillingOrm(id=2, order_id=1, first_name="bar")
order_orm = OrderOrm(id=1)
order_orm.billing.append(billing_orm_1)
order_orm.billing.append(billing_orm_2)

order_model = Order.from_orm(order_orm)
print(order_model) # 输出：id=1 name=['foo', 'bar']，完美符合预期

方法二：自定义GetterDict拦截属性访问 #

这种方法更适合复杂的映射场景，原理是自定义一个GetterDict，拦截Pydantic对ORM对象属性的获取操作，当请求name字段时直接返回处理后的值：

from pydantic import BaseModel
from typing import List
from sqlalchemy.orm import relationship
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from pydantic.utils import GetterDict

Base = declarative_base()

# ORM模型定义
class BillingOrm(Base):
    __tablename__ = "billing"
    id = Column(Integer, primary_key=True, nullable=False)
    order_id = Column(ForeignKey("orders.id", ondelete="CASCADE"), nullable=False)
    first_name = Column(String(20))

class OrderOrm(Base):
    __tablename__ = "orders"
    id = Column(Integer, primary_key=True, nullable=False)
    billing = relationship("BillingOrm")

# 自定义GetterDict
class OrderGetterDict(GetterDict):
    def get(self, key, default=None):
        if key == "name":
            # 当获取name字段时，直接返回ORM对象中billing列表的first_name集合
            return [b.first_name for b in self._obj.billing]
        # 其他字段按默认逻辑处理
        return super().get(key, default)

# Pydantic模型
class Order(BaseModel):
    id: int
    name: List[str] = None
    class Config:
        orm_mode = True
        getter_dict = OrderGetterDict # 指定使用自定义的GetterDict

# 测试
billing_orm_1 = BillingOrm(id=1, order_id=1, first_name="foo")
billing_orm_2 = BillingOrm(id=2, order_id=1, first_name="bar")
order_orm = OrderOrm(id=1)
order_orm.billing.append(billing_orm_1)
order_orm.billing.append(billing_orm_2)

order_model = Order.from_orm(order_orm)
print(order_model) # 输出：id=1 name=['foo', 'bar']，符合预期

这两种方法都可以直接用于FastAPI的response_model中，不需要修改接口的业务代码，非常方便。

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