apispec_oneofschema_polymorphic_marshmallow_openapi_documentation.py

1import marshmallow as mm
2from marshmallow_oneofschema import OneOfSchema
3from apispec import APISpec
4from apispec.ext.marshmallow import MarshmallowPlugin
5from apispec_oneofschema import MarshmallowPlugin as OneOfSchemaPlugin
6
7class CatSchema(mm.Schema):
8    name = mm.fields.String(required=True)
9    meow = mm.fields.Boolean(required=True)
10
11class DogSchema(mm.Schema):
12    name = mm.fields.String(required=True)
13    bark = mm.fields.Boolean(required=True)
14
15class PetSchema(OneOfSchema):
16    type_schemas = {"cat": CatSchema, "dog": DogSchema}
17
18    def get_obj_type(self, obj):
19        if isinstance(obj, Cat):
20            return "cat"
21        elif isinstance(obj, Dog):
22            return "dog"
23        else:
24            raise Exception("Unknown object type: {}".format(obj))
25
26# Create spec
27spec = APISpec(
28    title="Pet Store",
29    version="1.0.0",
30    openapi_version="3.0.2",
31    plugins=[MarshmallowPlugin(), OneOfSchemaPlugin()],
32)
33
34# Register the schema
35spec.components.schema("Pet", schema=PetSchema)
36
37# The generated spec will now use 'oneOf' for the Pet schema
38import json
39print(json.dumps(spec.to_dict(), indent=2))