python - 如何使用 anchor 将 pydantic 模型导出到 yaml？

Question

我想将 Pydantic 模型导出到 YAML，但要避免重复值并改用引用( anchor +别名)。

这是一个例子:

from typing import List
from ruamel.yaml import YAML  # type: ignore
import yaml
from pydantic import BaseModel

class Author(BaseModel):
    id: str
    name: str
    age: int

class Book(BaseModel):
    id: str
    title: str
    author: Author

class Library(BaseModel):
    authors: List[Author]
    books: List[Book]


john_smith = Author(id="auth1", name="John Smith", age=42)

books = [
    Book(id="book1", title="Some title", author=john_smith),
    Book(id="book2", title="Another one", author=john_smith),
]

library = Library(authors=[john_smith], books=books)

print(yaml.dump(library.dict()))

我得到:

authors:
- age: 42
  id: auth1
  name: John Smith
books:
- author:
    age: 42
    id: auth1
    name: John Smith
  id: book1
  title: Some title
- author:
    age: 42
    id: auth1
    name: John Smith
  id: book2
  title: Another one

您可以看到每本书中的所有作者字段都是重复的。我想要一些使用 anchor 而不是复制所有信息的东西，如下所示:

authors:
- &auth1
  age: 42
  id: auth1
  name: John Smith
books:
- author: *auth1
  id: book1
  title: Some title
- author: *auth1
  id: book2
  title: Another one

我怎样才能做到这一点？

Answer 1

当你遍历一个嵌套的 Python 数据结构来转换它时，你必须处理具有自引用的可能性，否则如果数据是自引用的，您的代码将陷入无限循环。

ruamel.yaml(和标准库 json.dump() )处理的方式是保持集合对象的 id() 列表(你想要递归，所以不是像 int、float、str 这样的基元，如果是 id()已经在列表中，该集合对象的第一次出现作为 anchor ，其他出现作为别名，因此您不必再次递归到对象中(json.dump() 告诉您它不能倾倒这样的一个结构，但至少它不会挂起)。

在 ruamel.yaml 中使用了相同的机制(跟踪 id()s)来避免重复相同的集合在多个其他集合中引用。

pydantic 似乎没有这样做，因此在调用 library.dict() 时会得到“写出”的结构。我认为这就是为什么在文档中告诉您使用字符串的原因有一个类名时 dumping pydanctic to JSON with self referential data

要绕过 pydantic 的这个限制，您可以做两件事:

• 编写一个 .dict() 的替代方案，它返回一个数据结构，该数据结构转储到所需的 YAML 文档格式，这意味着它需要在多个地方返回具有相同数据(dict)的结构。

• 确保您可以使用 ruamel.yaml 直接转储您的类，这样您就不必转换它们。

但是要使这两个都起作用，需要您添加到 book1 和 book2 的作者在添加后是相同的，但事实并非如此。你不能安全地假设如果两个字典有相同的键/值对，它们就是同一个对象所以任何比较都需要使用 is 而不是使用 ==。

将 john_smith 传递给 Book() 的两个调用后，您没有属性 .author指向相同的数据(即具有相同的 id()):

from pydantic import BaseModel
from typing import List

class Author(BaseModel):
    id: str
    name: str
    age: int

class Book(BaseModel):
    id: str
    title: str
    author: Author

class Library(BaseModel):
    authors: List[Author]
    books: List[Book]


john_smith = Author(id="auth1", name="John Smith", age=42)

books = [
    Book(id="book1", title="Some title", author=john_smith),
    Book(id="book2", title="Another one", author=john_smith),
]

library = Library(authors=[john_smith], books=books)

print('same author?',  john_smith is library.books[0].author)
print('same author?',  library.books[0].author is library.books[1].author)

给出:

same author? False
same author? False

你可以做的是强制作者相同，然后使用比 pydantic 的 .dict() 更聪明的东西:

import sys
import ruamel.yaml


def gen_data(d, id_map=None):
    if id_map is None:
        id_map = {}
    d_id = id(d)
    if d_id in id_map:
        print('already found', id_map)
        return id_map[d_id]
    if isinstance(d, BaseModel):
        ret_val = {}
        for k, v in d:
            if k == 'author':
                print('auth', v, id(v))
            ret_val[k] = gen_data(v, id_map)
    elif isinstance(d, list):
        ret_val = []
        for elem in d:
            ret_val.append(gen_data(elem, id_map))
    else:
        return d  # should be primitive
    id_map[d_id] = ret_val
    return ret_val

# force authors to be the same
library.books[0].author = library.books[1].author = library.authors[0]
assert  library.books[0].author is library.books[1].author

# alternative for .dict()
data = gen_data(library)
    
yaml = ruamel.yaml.YAML()
yaml.dump(data, sys.stdout)

这会产生您想要的结果:

auth id='auth1' name='John Smith' age=42 140494566559168
already found {140494566559168: {'id': 'auth1', 'name': 'John Smith', 'age': 42}, 140494576359168: [{'id': 'auth1', 'name': 'John Smith', 'age': 42}]}
auth id='auth1' name='John Smith' age=42 140494566559168
already found {140494566559168: {'id': 'auth1', 'name': 'John Smith', 'age': 42}, 140494576359168: [{'id': 'auth1', 'name': 'John Smith', 'age': 42}], 140494566559216: {'id': 'book1', 'title': 'Some title', 'author': {'id': 'auth1', 'name': 'John Smith', 'age': 42}}}
authors:
- &id001
  id: auth1
  name: John Smith
  age: 42
books:
- id: book1
  title: Some title
  author: *id001
- id: book2
  title: Another one
  author: *id001

请注意，您不应导入 yaml，而是实例化一个 ruamel.yaml.YAML() 实例。

如有必要，可以在 ruamel.yaml 中控制 anchor /别名的名称id001 以外的东西。