marshmallow-polyfieldを使ってoneOf的な構造のdataを扱う

はじめに

例えば、以下のよう１つのfieldに複数の形状の値が入ることがある。そして、その形状を決めるためにtypeなどfieldを含まれているJSONがあるとする。 以下の様な感じ(下の例では、personとgroupという2つの形状がobに入る可能性がある)。

{
  "ob": {
    "age": 20,
    "name": "foo"
  },
  "type": "person"
}

もしくはこう。

{
  "ob": {
    "name": "A",
    "members": [
      {
        "age": 20,
        "name": "foo"
      }
    ]
  },
  "type": "group"
}

それぞれ、typeで判別できるけれど。これを良い感じにmarshmallowでserialize,deserializeしたいという話し。

準備

事前に以下が必要。marshmallow-polyfieldを使う。

$ pip install marshmallow-polyfield

方法

以下の様な感じ。

from marshmallow import Schema, fields
from marshmallow_polyfield import PolyField


class Person(Schema):
    name = fields.String(required=True)
    age = fields.Integer(required=True)


class Group(Schema):
    name = fields.String(required=True)
    members = fields.List(fields.Nested(Person()), required=True)


def selector_for_deserialize(d, parent):
    if parent.get("type") == "group":
        return Group()
    else:
        return Person()


def selector_for_serialize(ob, parent):
    if "members" in ob:
        parent["type"] = "group"
        return Group()
    else:
        parent["type"] = "person"
        return Person()


class S(Schema):
    type = fields.String(required=True)
    ob = PolyField(
        serialization_schema_selector=selector_for_serialize,
        deserialization_schema_selector=selector_for_deserialize,
        required=True
    )


print(S().load({"ob": {"name": "foo", "age": 20}, "type": "person"}))
print(S().load({"ob": {"name": "A", "members": [{"name": "foo", "age": 20}]}, "type": "group"}))
print(S().dump({"ob": {"name": "foo", "age": 20}}))
print(S().dump({"ob": {"name": "A", "members": [{"name": "foo", "age": 20}]}}))

# UnmarshalResult(data={'ob': {'name': 'foo', 'age': 20}, 'type': 'person'}, errors={})
# UnmarshalResult(data={'ob': {'name': 'A', 'members': [{'name': 'foo', 'age': 20}]}, 'type': 'group'}, errors={})
# MarshalResult(data={'ob': {'name': 'foo', 'age': 20}, 'type': 'person'}, errors={})
# MarshalResult(data={'ob': {'name': 'A', 'members': [{'name': 'foo', 'age': 20}]}, 'type': 'group'}, errors={})