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

Question

我正在使用 JSON 架构中的 datamodel-code-generator to generate pydantic models。

这里是JSON schema使用的。

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

# 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

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

# 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']))

但是，我还没有找到让全局 class 使用 table 的扩展定义的方法。另外，简单地向 subclass 添加新字段似乎不起作用。

有没有一种方法可以扩展这些 classes 而不必修改 pydantic 生成的模型？或者也许是另一个从 JSON 模式生成 Python 代码的工具？

Answer 1

I haven't found a way to make the Globals class use the extended definition of the table

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

Also, simply adding new fields to the subclass does not seem to work

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

此示例展示了一种将计算字段 records_by_id 添加到 ExtendedTable 并在 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')}))

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

Is it possible to extend pydantic models generated using datamodel-code-generator?

python

pydantic