CODEWITHSUNDEEP
pythonpydantic
Ben
Ben

Reputation: 21615

How can I transform my input data to fit my desired Pydantic model

Suppose I have the following input data regarding a pet owner.

from types import SimpleNamespace

petowner1 = SimpleNamespace(
    id = 1,
    cats = [
            SimpleNamespace(id=1, name='Princess Peach')
        ],
    dogs = [
        SimpleNamespace(id=1, name='Sparky'),
        SimpleNamespace(id=2, name='Clifford')
    ]
)

petowner1 has an id, a list of cats, and a list of dogs. ...but I think it makes more sense for an Owner to have a list of Pets, each with a type attribute ('cat' or 'dog'). Hence, I set up the following Pydantic models

class Pet(BaseModel):
    id: int
    type: str
    name: str
    
    class Config:
        orm_mode = True
        

class Owner(BaseModel):
    id: int
    pets: List[Pet]

    class Config:
        orm_mode = True

Given my input data, how can I populate these Pydantic models? My end goal is to do something like

owner = Owner.from_orm(petowner1)
owner.json()

which should output

{
    'id': 1,
    'pets': [
        {'id': 1, 'type': 'cat', 'name': 'Princess Peach'},
        {'id': 1, 'type': 'dog', 'name': 'Sparky'},
        {'id': 2, 'type': 'dog', 'name': 'Clifford'}
    ]
}

Upvotes: 1

Views: 3526

Answers (2)

Ben
Ben

Reputation: 21615

Figured out a solution that involves subclassing GetterDict and extending its get() method, as described here.

from typing import Optional, List, Any
from pydantic import BaseModel
from pydantic.utils import GetterDict

class Pet(BaseModel):
    id: int
    pet_type: str
    name: str

    class Config:
        orm_mode = True


class OwnerGetter(GetterDict):

    def get(self, key: str, default: Any = None) -> Any:
        if key == 'pets':
            pets = [
                *[Pet(pet_type="cat", id=x.id, name=x.name) for x in self._obj.cats],
                *[Pet(pet_type="cat", id=x.id, name=x.name) for x in self._obj.dogs],
            ]
            return pets

        else:
            return super(OwnerGetter, self).get(key, default)


class Owner(BaseModel):
    id: int
    pets: List[Pet]

    class Config:
        orm_mode = True
        getter_dict = OwnerGetter

Usage

owner = Owner.from_orm(petowner1)
owner.json()

{
  "id": 1,
  "pets": [
    {
      "id": 1,
      "pet_type": "cat",
      "name": "Princess Peach"
    },
    {
      "id": 1,
      "pet_type": "cat",
      "name": "Sparky"
    },
    {
      "id": 2,
      "pet_type": "cat",
      "name": "Clifford"
    }
  ]
}

The caveat here is that we have to initialize each Pet using Pet(pet_type="cat", id=x.id, name=x.name) instead of Pet.from_orm(...). I've created a lengthier version of this with a custom PetGetter so that I can use

pets = [
  *[Pet.from_orm(SimpleNamespace(pet_type="cat", data=x)) for x in self._obj.cats],
  *[Pet.from_orm(SimpleNamespace(pet_type="dog", data=x)) for x in self._obj.dogs]
]

in place of

pets = [
  *[Pet(pet_type="cat", id=x.id, name=x.name) for x in self._obj.cats],
  *[Pet(pet_type="cat", id=x.id, name=x.name) for x in self._obj.dogs]
]

Upvotes: 2

kosciej16
kosciej16

Reputation: 7128

You can use simple loop to contruct object:

class Owner(BaseModel):
    id: int
    pets: List[Pet]

    class Config:
        orm_mode = True

    @classmethod
    def from_orm(cls, p):
        pets = [
            *[Pet(id=cat.id, type="cat", name=cat.name) for cat in p.cats],
            *[Pet(id=dog.id, type="dog", name=dog.name) for dog in p.dogs],
        ]
        return cls(id=p.id, pets=pets)

Upvotes: 2

Related Questions