如何转换我的输入数据以适合我想要的 Pydantic 模型

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

假设我有以下关于宠物主人的输入数据。

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 有一个 id、一个 cats 列表和一个 dogs 列表。 ...但我认为所有者拥有一个宠物列表更有意义,每个宠物都有一个type属性( 'cat' 或 'dog')。因此,我设置了以下 Pydantic 模型

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

根据我的输入数据,如何填充这些 Pydantic 模型?我的最终目标是做类似

的事情
owner = Owner.from_orm(petowner1)
owner.json()

应该输出

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

您可以使用简单的循环构造对象:

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)

想出一个解决方案,涉及子类化 GetterDict 和扩展其 get() 方法,如 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

用法

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"
    }
  ]
}

这里需要注意的是,我们必须使用 Pet(pet_type="cat", id=x.id, name=x.name) 而不是 Pet.from_orm(...) 来初始化每个 Pet。我用自定义 PetGetter 创建了一个更长的版本,这样我就可以使用

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]
]

代替

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]
]