是否可以扩展使用DataModel-Code-Menerator生成的Pydantic模型？

Question

我正在使用 datamodel-code-generator 生成 pydantic Models 来自JSON模式。

这是 json schema> json Schema 使用。

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

# 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 __ init __（）方法中的实例属性，但是将字段声明为类属性。

此示例显示了一种将计算出的字段Records_by_id添加到ExtendedTable并使用ExtendedTable in extendedGlobals：

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