marshmallow_oneofschema_polymorphic_serialization_with_type_marker.py

1from marshmallow import Schema, fields
2from marshmallow_oneofschema import OneOfSchema
3
4
5class Elephant:
6    def __init__(self, name, size):
7        self.name = name
8        self.size = size
9
10
11class Zebra:
12    def __init__(self, name, stripes):
13        self.name = name
14        self.stripes = stripes
15
16
17class ElephantSchema(Schema):
18    name = fields.Str()
19    size = fields.Str()
20
21
22class ZebraSchema(Schema):
23    name = fields.Str()
24    stripes = fields.Int()
25
26
27class ZooAnimalSchema(OneOfSchema):
28    type_schemas = {"elephant": ElephantSchema, "zebra": ZebraSchema}
29
30    def get_obj_type(self, obj):
31        if isinstance(obj, Elephant):
32            return "elephant"
33        elif isinstance(obj, Zebra):
34            return "zebra"
35        else:
36            raise Exception("Unknown object type: {}".format(obj))
37
38
39# Serialization
40zoo_data = ZooAnimalSchema(many=True).dump([
41    Elephant(name="Elephant 1", size="large"),
42    Zebra(name="Zebra 1", stripes=10),
43])
44# [
45#    {"name": "Elephant 1", "size": "large", "type": "elephant"},
46#    {"name": "Zebra 1", "stripes": 10, "type": "zebra"}
47# ]
48
49# Deserialization
50zoo_animals = ZooAnimalSchema(many=True).load([
51    {"name": "Elephant 2", "size": "small", "type": "elephant"},
52    {"name": "Zebra 2", "stripes": 20, "type": "zebra"},
53])
54# [
55#    {'name': 'Elephant 2', 'size': 'small'},
56#    {'name': 'Zebra 2', 'stripes': 20}
57# ]