Skip to content

Commit a4d57ee

Browse files
committed
Add support for marshmallow-oneofschema polymorphic schemas (#1012)
Support polymorphic schemas defined with marshmallow-oneofschema.OneOfSchema in MarshmallowPlugin: - Check for marshmallow-oneofschema availability at import time in OpenAPIConverter. - Convert OneOfSchema instances and classes to OpenAPI 3 oneOf schemas with a discriminator object mapping type names to component schema references. - Preserve schema title and description from class Meta when present. - Support nested OneOfSchema instances and anonymous/inline schemas. - Add unit and integration tests covering OneOfSchema conversion and OpenAPI validation. - Add marshmallow-oneofschema as an optional test dependency in pyproject.toml and uv.lock.
1 parent 0b66a26 commit a4d57ee

6 files changed

Lines changed: 340 additions & 1 deletion

File tree

‎CHANGELOG.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@ Changelog
44
unreleased
55
**********
66

7+
Features:
8+
9+
- ``MarshmallowPlugin``: Add support for polymorphic schemas using ``marshmallow-oneofschema.OneOfSchema`` (:issue:`1012`).
10+
711
Bug fixes:
812

913
- Recognize ``trace`` operations defined in view docstrings, instead of

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ marshmallow = ["marshmallow>=3.18.0"]
4646

4747
[dependency-groups]
4848
docs = ["marshmallow>=4.0.0", "pyyaml", "sphinx-issues", "furo", "sphinx>=8.1"]
49-
tests = ["openapi-spec-validator", "pytest"]
49+
tests = ["marshmallow-oneofschema", "openapi-spec-validator", "pytest"]
5050
lint = ["pre-commit>=3.5,<5.0"]
5151
dev = [
5252
{ include-group = "tests" },

‎src/apispec/ext/marshmallow/openapi.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,13 @@
1919
from apispec import APISpec
2020
from apispec.exceptions import APISpecError
2121

22+
try:
23+
from marshmallow_oneofschema import OneOfSchema
24+
25+
ONEOFSCHEMA_AVAILABLE = True
26+
except ImportError:
27+
ONEOFSCHEMA_AVAILABLE = False
28+
2229
from .common import (
2330
get_fields,
2431
get_unique_schema_name,
@@ -248,6 +255,12 @@ def schema2jsonschema(self, schema):
248255
:param Schema schema: A marshmallow Schema instance
249256
:rtype: dict, a JSON Schema Object
250257
"""
258+
if ONEOFSCHEMA_AVAILABLE and (
259+
isinstance(schema, OneOfSchema)
260+
or (isinstance(schema, type) and issubclass(schema, OneOfSchema))
261+
):
262+
return self._oneof_schema2jsonschema(schema)
263+
251264
fields = get_fields(schema)
252265
Meta = getattr(schema, "Meta", None)
253266
partial = getattr(schema, "partial", None)
@@ -267,6 +280,60 @@ def schema2jsonschema(self, schema):
267280

268281
return jsonschema
269282

283+
def _oneof_schema2jsonschema(self, schema):
284+
"""Handle OneOfSchema instances by generating oneOf with discriminator.
285+
286+
:param OneOfSchema schema: A OneOfSchema instance or class
287+
:rtype: dict, a JSON Schema Object with oneOf
288+
"""
289+
if self.openapi_version.major < 3:
290+
raise APISpecError(
291+
f"OneOfSchema is only supported in OpenAPI 3.0.0 and above. "
292+
f"Current OpenAPI version is {self.openapi_version}."
293+
)
294+
295+
schema_instance = resolve_schema_instance(schema)
296+
type_schemas = getattr(schema_instance, "type_schemas", {})
297+
type_field = getattr(schema_instance, "type_field", "type")
298+
299+
if not type_schemas:
300+
return {"type": "object", "properties": {}}
301+
302+
one_of_list = []
303+
mapping = {}
304+
for type_name, type_schema in type_schemas.items():
305+
schema_ref = self.resolve_nested_schema(type_schema)
306+
one_of_list.append(schema_ref)
307+
308+
if isinstance(schema_ref, dict):
309+
if "$ref" in schema_ref:
310+
mapping[type_name] = schema_ref["$ref"]
311+
elif (
312+
"items" in schema_ref
313+
and isinstance(schema_ref["items"], dict)
314+
and "$ref" in schema_ref["items"]
315+
):
316+
mapping[type_name] = schema_ref["items"]["$ref"]
317+
elif isinstance(schema_ref, str):
318+
mapping[type_name] = schema_ref
319+
320+
discriminator = {"propertyName": type_field}
321+
if mapping:
322+
discriminator["mapping"] = mapping
323+
324+
jsonschema = {
325+
"oneOf": one_of_list,
326+
"discriminator": discriminator,
327+
}
328+
329+
Meta = getattr(schema, "Meta", getattr(schema_instance, "Meta", None))
330+
if hasattr(Meta, "title"):
331+
jsonschema["title"] = Meta.title
332+
if hasattr(Meta, "description"):
333+
jsonschema["description"] = Meta.description
334+
335+
return jsonschema
336+
270337
def fields2jsonschema(self, fields, *, partial=None):
271338
"""Return the JSON Schema Object given a mapping between field names and
272339
:class:`Field <marshmallow.Field>` objects.

‎tests/test_ext_marshmallow.py‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
get_paths,
3737
get_responses,
3838
get_schemas,
39+
validate_spec,
3940
)
4041

4142
MA_VERSION = Version(importlib.metadata.version("marshmallow"))
@@ -1369,3 +1370,146 @@ class SchemaWithTimeDelta(Schema):
13691370
get_schemas(spec)["SchemaWithTimeDelta"]["properties"]["day"]["x-unit"]
13701371
== "days"
13711372
)
1373+
1374+
1375+
try:
1376+
from marshmallow_oneofschema import OneOfSchema
1377+
1378+
ONEOFSCHEMA_AVAILABLE = True
1379+
except ImportError:
1380+
ONEOFSCHEMA_AVAILABLE = False
1381+
1382+
1383+
if ONEOFSCHEMA_AVAILABLE:
1384+
1385+
class OneOfFooSchema(Schema):
1386+
foo = String()
1387+
1388+
class OneOfBarSchema(Schema):
1389+
bar = Int()
1390+
1391+
class ExampleOneOf(OneOfSchema):
1392+
type_field = "type"
1393+
type_schemas = {
1394+
"foo": OneOfFooSchema,
1395+
"bar": OneOfBarSchema,
1396+
}
1397+
1398+
class ExampleParentSchema(Schema):
1399+
item = Nested(ExampleOneOf)
1400+
1401+
1402+
@pytest.mark.skipif(
1403+
not ONEOFSCHEMA_AVAILABLE, reason="marshmallow-oneofschema not installed"
1404+
)
1405+
class TestOneOfSchemaIntegration:
1406+
@pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
1407+
def test_spec_components_schema_with_oneof(self, openapi_version):
1408+
spec = APISpec(
1409+
title="OneOf Test",
1410+
version="1.0.0",
1411+
openapi_version=openapi_version,
1412+
plugins=[MarshmallowPlugin()],
1413+
)
1414+
spec.components.schema("ExampleOneOf", schema=ExampleOneOf)
1415+
1416+
schemas = get_schemas(spec)
1417+
assert "ExampleOneOf" in schemas
1418+
assert "OneOfFoo" in schemas
1419+
assert "OneOfBar" in schemas
1420+
assert schemas["ExampleOneOf"] == {
1421+
"oneOf": [
1422+
build_ref(spec, "schema", "OneOfFoo"),
1423+
build_ref(spec, "schema", "OneOfBar"),
1424+
],
1425+
"discriminator": {
1426+
"propertyName": "type",
1427+
"mapping": {
1428+
"foo": build_ref(spec, "schema", "OneOfFoo")["$ref"],
1429+
"bar": build_ref(spec, "schema", "OneOfBar")["$ref"],
1430+
},
1431+
},
1432+
}
1433+
validate_spec(spec)
1434+
1435+
@pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
1436+
def test_spec_nested_oneof_schema(self, openapi_version):
1437+
spec = APISpec(
1438+
title="OneOf Nested Test",
1439+
version="1.0.0",
1440+
openapi_version=openapi_version,
1441+
plugins=[MarshmallowPlugin()],
1442+
)
1443+
spec.components.schema("Parent", schema=ExampleParentSchema)
1444+
1445+
schemas = get_schemas(spec)
1446+
assert "Parent" in schemas
1447+
assert "ExampleOneOf" in schemas
1448+
assert "OneOfFoo" in schemas
1449+
assert "OneOfBar" in schemas
1450+
assert schemas["Parent"]["properties"]["item"] == build_ref(
1451+
spec, "schema", "ExampleOneOf"
1452+
)
1453+
validate_spec(spec)
1454+
1455+
@pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
1456+
def test_spec_operations_with_oneof(self, openapi_version):
1457+
spec = APISpec(
1458+
title="OneOf Operations Test",
1459+
version="1.0.0",
1460+
openapi_version=openapi_version,
1461+
plugins=[MarshmallowPlugin()],
1462+
)
1463+
spec.path(
1464+
"/items",
1465+
operations={
1466+
"post": {
1467+
"requestBody": {
1468+
"content": {"application/json": {"schema": ExampleOneOf}}
1469+
},
1470+
"responses": {
1471+
"200": {
1472+
"description": "success",
1473+
"content": {"application/json": {"schema": ExampleOneOf}},
1474+
}
1475+
},
1476+
}
1477+
},
1478+
)
1479+
schemas = get_schemas(spec)
1480+
assert "ExampleOneOf" in schemas
1481+
assert "OneOfFoo" in schemas
1482+
assert "OneOfBar" in schemas
1483+
paths = get_paths(spec)
1484+
assert paths["/items"]["post"]["requestBody"]["content"]["application/json"][
1485+
"schema"
1486+
] == build_ref(spec, "schema", "ExampleOneOf")
1487+
validate_spec(spec)
1488+
1489+
@pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
1490+
def test_spec_anonymous_oneof_inline(self, openapi_version):
1491+
spec = APISpec(
1492+
title="Anonymous OneOf",
1493+
version="1.0.0",
1494+
openapi_version=openapi_version,
1495+
plugins=[MarshmallowPlugin(schema_name_resolver=lambda s: None)],
1496+
)
1497+
spec.path(
1498+
"/items",
1499+
operations={
1500+
"post": {
1501+
"requestBody": {
1502+
"content": {"application/json": {"schema": ExampleOneOf}}
1503+
},
1504+
"responses": {"200": {"description": "success"}},
1505+
}
1506+
},
1507+
)
1508+
paths = get_paths(spec)
1509+
schema = paths["/items"]["post"]["requestBody"]["content"]["application/json"][
1510+
"schema"
1511+
]
1512+
assert "oneOf" in schema
1513+
assert len(schema["oneOf"]) == 2
1514+
assert schema["discriminator"] == {"propertyName": "type"}
1515+
validate_spec(spec)

‎tests/test_ext_marshmallow_openapi.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -717,3 +717,111 @@ def test_properties(self, field, properties, spec):
717717
for attr, expected_value in properties.items():
718718
assert attr in result
719719
assert result[attr] == expected_value
720+
721+
722+
try:
723+
from marshmallow_oneofschema import OneOfSchema
724+
725+
ONEOFSCHEMA_AVAILABLE = True
726+
except ImportError:
727+
ONEOFSCHEMA_AVAILABLE = False
728+
729+
730+
if ONEOFSCHEMA_AVAILABLE:
731+
732+
class FooSchema(Schema):
733+
foo = fields.Str()
734+
735+
class BarSchema(Schema):
736+
bar = fields.Int()
737+
738+
class MyOneOf(OneOfSchema):
739+
type_field = "type"
740+
type_schemas = {
741+
"foo": FooSchema,
742+
"bar": BarSchema,
743+
}
744+
745+
class CustomFieldOneOf(OneOfSchema):
746+
type_field = "custom_type"
747+
type_schemas = {
748+
"foo": FooSchema,
749+
"bar": BarSchema,
750+
}
751+
752+
class TitledOneOf(OneOfSchema):
753+
type_field = "type"
754+
type_schemas = {
755+
"foo": FooSchema,
756+
"bar": BarSchema,
757+
}
758+
759+
class Meta:
760+
title = "PolymorphicOffer"
761+
description = "A polymorphic offer"
762+
763+
class EmptyOneOf(OneOfSchema):
764+
type_schemas: dict[str, type[Schema]] = {}
765+
766+
767+
@pytest.mark.skipif(
768+
not ONEOFSCHEMA_AVAILABLE, reason="marshmallow-oneofschema not installed"
769+
)
770+
class TestOneOfSchemaToOpenAPI:
771+
@pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
772+
def test_oneof_schema_to_jsonschema(self, openapi):
773+
res = openapi.schema2jsonschema(MyOneOf)
774+
assert res == {
775+
"oneOf": [
776+
{"$ref": "#/components/schemas/Foo"},
777+
{"$ref": "#/components/schemas/Bar"},
778+
],
779+
"discriminator": {
780+
"propertyName": "type",
781+
"mapping": {
782+
"foo": "#/components/schemas/Foo",
783+
"bar": "#/components/schemas/Bar",
784+
},
785+
},
786+
}
787+
788+
@pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
789+
def test_oneof_schema_instance_to_jsonschema(self, openapi):
790+
res = openapi.schema2jsonschema(MyOneOf())
791+
assert res == {
792+
"oneOf": [
793+
{"$ref": "#/components/schemas/Foo"},
794+
{"$ref": "#/components/schemas/Bar"},
795+
],
796+
"discriminator": {
797+
"propertyName": "type",
798+
"mapping": {
799+
"foo": "#/components/schemas/Foo",
800+
"bar": "#/components/schemas/Bar",
801+
},
802+
},
803+
}
804+
805+
@pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
806+
def test_oneof_custom_type_field(self, openapi):
807+
res = openapi.schema2jsonschema(CustomFieldOneOf)
808+
assert res["discriminator"]["propertyName"] == "custom_type"
809+
810+
@pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
811+
def test_oneof_meta_title_and_description(self, openapi):
812+
res = openapi.schema2jsonschema(TitledOneOf)
813+
assert res["title"] == "PolymorphicOffer"
814+
assert res["description"] == "A polymorphic offer"
815+
816+
@pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
817+
def test_oneof_empty_type_schemas(self, openapi):
818+
res = openapi.schema2jsonschema(EmptyOneOf)
819+
assert res == {"type": "object", "properties": {}}
820+
821+
@pytest.mark.parametrize("openapi", ("2.0",), indirect=True)
822+
def test_oneof_openapi_v2_raises_error(self, openapi):
823+
with pytest.raises(
824+
exceptions.APISpecError,
825+
match="OneOfSchema is only supported in OpenAPI 3.0.0 and above",
826+
):
827+
openapi.schema2jsonschema(MyOneOf)

0 commit comments

Comments
 (0)