是否可以扩展使用数据模型代码生成器生成的pydantic模型?

Question

我正在使用数据模型代码生成器从JSON模式生成pydantic模型。

下面是使用的JSON模式。

以及运行数据模型代码生成器后生成的模型。

# File: datamodel.py
from __future__ import annotations
from typing import List
from pydantic import BaseModel
class Record(BaseModel):
id: int
name: str
class Table(BaseModel):
records: List[Record]
class Globals(BaseModel):
table: Table

我一直在尝试用新属性扩展生成的类。

# File: extensions.py
import json
from datamodel import Table, Globals
class ExtendedTable(Table):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
print('ExtendedTable Constructor')
# Won't work because "ExtendedTable" object has no field "records_by_id"
self.records_by_id = {record.id: record for record in self.records}
class ExtendedGlobals(Globals):
def __init__(self, table: ExtendedTable):
super().__init__(table=table)
print('ExtendedGlobals Constructor')
if __name__ == '__main__':
records = '''
{
"table": {
"records": [{"id": 0, "name": "A"}, {"id": 1, "name": "B"}]
}
}
'''
content = json.loads(records)
# Both won't call ExtendedTable.__init__()
ExtendedGlobals(**content)
ExtendedGlobals.parse_obj(content)
ExtendedGlobals(table=ExtendedTable(**content['table']))

然而，我还没有找到一种方法来使Globals类使用表的扩展定义。此外，简单地向子类中添加新字段似乎也不起作用。

是否有一种方法来扩展这些类而不必修改pydantic生成的模型?或者是另一个从JSON模式生成Python代码的工具?

Answer 1

我还没有找到一种方法来使Globals类使用表的扩展定义

如果您再次使用所需的类型声明字段，则可以更改子类中字段的类型。

此外，简单地向子类添加新字段似乎不起作用

看起来你是在__init__()方法中设置实例属性，但是字段被声明为类属性。

这个例子展示了在ExtendedTable中添加计算字段records_by_id，在ExtendedGlobals中使用ExtendedTable的方法:

# File: extensions.py
import json
from typing import Any, Dict, List, Optional
from pydantic import Field, validator
from datamodel import Globals, Record, Table

class ExtendedTable(Table):
# New fields are declared as class attributes not as instance attributes inside the __init__()
# Calculated fields usually have a default value or default factory so that you don't have to provide a value
# I prefer a default_factory for mutable values
records_by_id: Dict[int, Record] = Field(default_factory=dict)
# A validator can populate a calculated field
# Use always=True to run the validator even if a value is not supplied and the default value is used
@validator("records_by_id", always=True)
def _calculate_records_by_id(
cls, value: Dict[int, Record], values: Dict[str, Any]
) -> Dict[int, Record]:
records: Optional[List[Record]] = values.get("records")
if records is None:
# The records field was not valid
# Return value or raise a ValueError instead if you want
return value
return {record.id: record for record in records}

class ExtendedGlobals(Globals):
# You can change the type of a field in a subclass if you declare the field again
table: ExtendedTable

if __name__ == "__main__":
records = """
{
"table": {
"records": [{"id": 0, "name": "A"}, {"id": 1, "name": "B"}]
}
}
"""
content = json.loads(records)
extended_globals = ExtendedGlobals.parse_obj(content)
print(repr(extended_globals))

输出:

ExtendedGlobals(table=ExtendedTable(records=[Record(id=0, name='A'), Record(id=1, name='B')], records_by_id={0: Record(id=0, name='A'), 1: Record(id=1, name='B')}))