Best way to flatten and remap ORM to Pydantic Model

Question:

I am using Pydantic with FastApi to output ORM data into JSON. I would like to flatten and remap the ORM model to eliminate an unnecessary level in the JSON.

Here’s a simplified example to illustrate the problem.

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

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

How to map values from nested dict to Pydantic Model? provides a solution that works for dictionaries by using the init function in the Pydantic model class. This example shows how that works with dictionaries:

from pydantic import BaseModel

# The following approach works with a dictionary as the input

order_dict = {"id": 1, "billing": {"first_name": "foo"}}

# desired output: {"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 for dictionary: {'id': 1, 'billing': {'first_name': 'foo'}}
        kwargs["name"] = kwargs["billing"]["first_name"]
        super().__init__(**kwargs)


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

(This script is complete, it should run "as is")

However, when working with ORM objects, this approach does not work. It appears that the init function is not called. Here’s an example which will not provide the desired output.

from pydantic import BaseModel, root_validator
from typing import List
from sqlalchemy.orm import relationship
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.dialects.postgresql import ARRAY
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()

from pydantic.utils import GetterDict

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")


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]  # uncomment to verify the relationship is working

    class Config:
        orm_mode = True

    def __init__(self, **kwargs):
        # This __init__ function does not run when using from_orm to parse ORM object
        print("kwargs for orm:", kwargs)
        kwargs["name"] = kwargs["billing"]["first_name"]
        super().__init__(**kwargs)


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)
# Output returns 'None' for name instead of ['foo','bar']
print(order_model)  # id=1 name=None

(This script is complete, it should run "as is")

The output returns name=None instead of the desired list of names.

In the above example, I am using Order.from_orm to create the Pydantic model. This approach seems to be the same that is used by FastApi when specifying a response model. The desired solution should support use in the FastApi response model as shown in this example:

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

Update:
Regarding MatsLindh comment to try validators, I replaced the init function with a root validator, however, I’m unable to mutate the return values to include a new attribute. I suspect this issue is because it is a ORM object and not a true dictionary. The following code will extract the names and print them in the desired list. However, I can’t see how to include this updated result in the model response:

@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)
        # values["name"] = names # error: 'GetterDict' object does not support item assignment
        return values

I also found a couple other discussions on this problem that led me to try this approach:
https://github.com/samuelcolvin/pydantic/issues/717
https://gitmemory.com/issue/samuelcolvin/pydantic/821/744047672

Asked By: STEM FabLab

||

Answers:

What if you override the from_orm class method?

class Order(BaseModel):
    id: int
    name: List[str] = None
    billing: List[Billing]

    class Config:
        orm_mode = True

    @classmethod
    def from_orm(cls, obj: Any) -> 'Order':
        # `obj` is the orm model instance
        if hasattr(obj, 'billing'):
            obj.name = obj.billing.first_name
        return super().from_orm(obj)
Answered By: Nick Petrovic

I really missed the handy Django REST Framework serializers while working with the FastAPI + Pydantic stack… So I wrangled with GetterDict to allow defining field getter function in the Pydantic model like this:

class User(FromORM):

    fullname: str

    class Config(FromORM.Config):
        getter_dict = FieldGetter.bind(lambda: User)

    @staticmethod
    def get_fullname(obj: User) -> str:
        return f'{obj.firstname} {obj.lastname}'

where the magic part FieldGetter is implemented as

from typing import Any, Callable, Optional, Type
from types import new_class
from pydantic import BaseModel
from pydantic.utils import GetterDict


class FieldGetter(GetterDict):

    model_class_forward_ref: Optional[Callable] = None
    model_class: Optional[Type[BaseModel]] = None

    def __new__(cls, *args, **kwargs):
        inst = super().__new__(cls)
        if cls.model_class_forward_ref:
            inst.model_class = cls.model_class_forward_ref()

        return inst

    @classmethod
    def bind(cls, model_class_forward_ref: Callable):
        sub_class = new_class(f'{cls.__name__}FieldGetter', (cls,))
        sub_class.model_class_forward_ref = model_class_forward_ref
        return sub_class

    def get(self, key: str, default):
        if hasattr(self._obj, key):
            return super().get(key, default)

        getter_fun_name = f'get_{key}'
        if not (getter := getattr(self.model_class, getter_fun_name, None)):
            raise AttributeError(f'no field getter function found for {key}')

        return getter(self._obj)


class FromORM(BaseModel):

    class Config:
        orm_mode = True
        getter_dict = FieldGetter
Answered By: Namoshizun