使用 Pydantic 解析不同模型的列表

Question

我正在尝试创建一个 pydantic 模型，它可以用不同类型的模型列表解析 json 字符串。请参阅下面的示例

test_object = """
{
  "distributions":
    [
      {
        "param_name": "test1",
        "attributes":
          {
            "distribution_name": "UniformDistribution",
            "low": "1.0",
            "high": "2.0"
          }
      },
      {
        "param_name": "test2",
        "attributes":
          {
            "distribution_name": "UniformDistribution",
            "low": "1.0",
            "high":"2.0"
          }
      },
      {
        "param_name": "test3",
        "attributes":
          {
            "distribution_name": "IntUniformDistribution",
            "low": "1",
            "high": "2",
            "q": 4
          }
      },
      {
        "param_name": "test4",
        "attributes":
          {
            "distribution_name": "DiscreteUniformDistribution",
            "low": "1.0",
            "high": "2.0",
            "step": ".1"
          }
      }
    ]
}
"""

这是我用来解析上述字符串的 pydantic 模型。

from pydantic import BaseModel, conint, confloat, ValidationError
from typing import List, Literal, Union


class DiscreteUniformDistribution(BaseModel):
    distribution_name: Literal["DiscreteUniformDistribution"]
    low: float
    high: float
    q: confloat(gt=0)

class IntUniformDistribution(BaseModel):
    distribution_name: Literal["IntUniformDistribution"]
    low: int
    high: int
    step: conint(gt=0)

class UniformDistribution(BaseModel):
    distribution_name: Literal["UniformDistribution"]
    low: float
    high: float

class Distribution(BaseModel):
    param_name: str
    attributes: Union[
        DiscreteUniformDistribution,
        IntUniformDistribution,
        UniformDistribution,
    ]

class Distributions(BaseModel):
    distributions: List[Distribution]


try:
    Distributions.parse_raw(test_object)
except ValidationError as e:
    print(e)

但是我收到一个错误

9 validation errors for Distributions
distributions -> 2 -> attributes -> distribution_name
  unexpected value; permitted: 'DiscreteUniformDistribution' (type=value_error.const; given=IntUniformDistribution; permitted=('DiscreteUniformDistribution',))
distributions -> 2 -> attributes -> step
  field required (type=value_error.missing)
distributions -> 2 -> attributes -> distribution_name
  unexpected value; permitted: 'UniformDistribution' (type=value_error.const; given=IntUniformDistribution; permitted=('UniformDistribution',))
distributions -> 3 -> attributes -> q
  field required (type=value_error.missing)
distributions -> 3 -> attributes -> distribution_name
  unexpected value; permitted: 'IntUniformDistribution' (type=value_error.const; given=DiscreteUniformDistribution; permitted=('IntUniformDistribution',))
distributions -> 3 -> attributes -> low
  value is not a valid integer (type=type_error.integer)
distributions -> 3 -> attributes -> high
  value is not a valid integer (type=type_error.integer)
distributions -> 3 -> attributes -> step
  value is not a valid integer (type=type_error.integer)
distributions -> 3 -> attributes -> distribution_name
  unexpected value; permitted: 'UniformDistribution' (type=value_error.const; given=DiscreteUniformDistribution; permitted=('UniformDistribution',))

如果 distributions 列表中只有元素，则工作正常，如下所示。

test_object = """
{
  "distributions": 
    [
      {
        "param_name": "test1",
        "attributes":
          {
            "distribution_name": "UniformDistribution",
            "low": "1.0",
            "high": "2.0"
          }
      }
    ]
}
"""

我是 pydantic 的新手。我怀疑这个错误与我对 Union 的不当使用有关。

Answer 1

您对 Union[] 的用法看起来不错，但是，您的模型定义中有错字。

您需要将 DiscreteUniformDistribution() 中的 q 与 IntUniformDistribution() 中的 step 交换（仅字段名称，而不是类型），即：

class DiscreteUniformDistribution(BaseModel):
    distribution_name: Literal['DiscreteUniformDistribution']
    low: float
    high: float
    step: confloat(gt=0)
#   ^^^^
#   This is called q in your definition

class IntUniformDistribution(BaseModel):
    distribution_name: Literal['IntUniformDistribution']
    low: int
    high: int
    q: conint(gt=0)
#   ^
#   This is called step in your definition

使用 Pydantic 解析不同模型的列表

Parsing list of different models with Pydantic

python

pydantic