Add support for marshmallow-oneofschema polymorphic schemas (#1012)

CHANGELOG.rst

@@ -4,6 +4,10 @@ Changelog
 unreleased
 **********
 
+Features:
+
+- ``MarshmallowPlugin``: Add support for polymorphic schemas using ``marshmallow-oneofschema.OneOfSchema`` (:issue:`1012`).
+
 Bug fixes:
 
 - Recognize ``trace`` operations defined in view docstrings, instead of

pyproject.toml

@@ -46,7 +46,7 @@ marshmallow = ["marshmallow>=3.18.0"]
 
 [dependency-groups]
 docs = ["marshmallow>=4.0.0", "pyyaml", "sphinx-issues", "furo", "sphinx>=8.1"]
-tests = ["openapi-spec-validator", "pytest"]
+tests = ["marshmallow-oneofschema", "openapi-spec-validator", "pytest"]
 lint = ["pre-commit>=3.5,<5.0"]
 dev = [
   { include-group = "tests" },

src/apispec/ext/marshmallow/openapi.py

@@ -19,6 +19,13 @@
 from apispec import APISpec
 from apispec.exceptions import APISpecError
 
+try:
+    from marshmallow_oneofschema import OneOfSchema
+
+    ONEOFSCHEMA_AVAILABLE = True
+except ImportError:
+    ONEOFSCHEMA_AVAILABLE = False
+
 from .common import (
     get_fields,
     get_unique_schema_name,
@@ -248,6 +255,12 @@ def schema2jsonschema(self, schema):
         :param Schema schema: A marshmallow Schema instance
         :rtype: dict, a JSON Schema Object
         """
+        if ONEOFSCHEMA_AVAILABLE and (
+            isinstance(schema, OneOfSchema)
+            or (isinstance(schema, type) and issubclass(schema, OneOfSchema))
+        ):
+            return self._oneof_schema2jsonschema(schema)
+
         fields = get_fields(schema)
         Meta = getattr(schema, "Meta", None)
         partial = getattr(schema, "partial", None)
@@ -267,6 +280,60 @@ def schema2jsonschema(self, schema):
 
         return jsonschema
 
+    def _oneof_schema2jsonschema(self, schema):
+        """Handle OneOfSchema instances by generating oneOf with discriminator.
+
+        :param OneOfSchema schema: A OneOfSchema instance or class
+        :rtype: dict, a JSON Schema Object with oneOf
+        """
+        if self.openapi_version.major < 3:
+            raise APISpecError(
+                f"OneOfSchema is only supported in OpenAPI 3.0.0 and above. "
+                f"Current OpenAPI version is {self.openapi_version}."
+            )
+
+        schema_instance = resolve_schema_instance(schema)
+        type_schemas = getattr(schema_instance, "type_schemas", {})
+        type_field = getattr(schema_instance, "type_field", "type")
+
+        if not type_schemas:
+            return {"type": "object", "properties": {}}
+
+        one_of_list = []
+        mapping = {}
+        for type_name, type_schema in type_schemas.items():
+            schema_ref = self.resolve_nested_schema(type_schema)
+            one_of_list.append(schema_ref)
+
+            if isinstance(schema_ref, dict):
+                if "$ref" in schema_ref:
+                    mapping[type_name] = schema_ref["$ref"]
+                elif (
+                    "items" in schema_ref
+                    and isinstance(schema_ref["items"], dict)
+                    and "$ref" in schema_ref["items"]
+                ):
+                    mapping[type_name] = schema_ref["items"]["$ref"]
+            elif isinstance(schema_ref, str):
+                mapping[type_name] = schema_ref
+
+        discriminator = {"propertyName": type_field}
+        if mapping:
+            discriminator["mapping"] = mapping
+
+        jsonschema = {
+            "oneOf": one_of_list,
+            "discriminator": discriminator,
+        }
+
+        Meta = getattr(schema, "Meta", getattr(schema_instance, "Meta", None))
+        if hasattr(Meta, "title"):
+            jsonschema["title"] = Meta.title
+        if hasattr(Meta, "description"):
+            jsonschema["description"] = Meta.description
+
+        return jsonschema
+
     def fields2jsonschema(self, fields, *, partial=None):
         """Return the JSON Schema Object given a mapping between field names and
         :class:`Field <marshmallow.Field>` objects.

tests/test_ext_marshmallow.py

@@ -36,6 +36,7 @@
     get_paths,
     get_responses,
     get_schemas,
+    validate_spec,
 )
 
 MA_VERSION = Version(importlib.metadata.version("marshmallow"))
@@ -1369,3 +1370,146 @@ class SchemaWithTimeDelta(Schema):
             get_schemas(spec)["SchemaWithTimeDelta"]["properties"]["day"]["x-unit"]
             == "days"
         )
+
+
+try:
+    from marshmallow_oneofschema import OneOfSchema
+
+    ONEOFSCHEMA_AVAILABLE = True
+except ImportError:
+    ONEOFSCHEMA_AVAILABLE = False
+
+
+if ONEOFSCHEMA_AVAILABLE:
+
+    class OneOfFooSchema(Schema):
+        foo = String()
+
+    class OneOfBarSchema(Schema):
+        bar = Int()
+
+    class ExampleOneOf(OneOfSchema):
+        type_field = "type"
+        type_schemas = {
+            "foo": OneOfFooSchema,
+            "bar": OneOfBarSchema,
+        }
+
+    class ExampleParentSchema(Schema):
+        item = Nested(ExampleOneOf)
+
+
+@pytest.mark.skipif(
+    not ONEOFSCHEMA_AVAILABLE, reason="marshmallow-oneofschema not installed"
+)
+class TestOneOfSchemaIntegration:
+    @pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
+    def test_spec_components_schema_with_oneof(self, openapi_version):
+        spec = APISpec(
+            title="OneOf Test",
+            version="1.0.0",
+            openapi_version=openapi_version,
+            plugins=[MarshmallowPlugin()],
+        )
+        spec.components.schema("ExampleOneOf", schema=ExampleOneOf)
+
+        schemas = get_schemas(spec)
+        assert "ExampleOneOf" in schemas
+        assert "OneOfFoo" in schemas
+        assert "OneOfBar" in schemas
+        assert schemas["ExampleOneOf"] == {
+            "oneOf": [
+                build_ref(spec, "schema", "OneOfFoo"),
+                build_ref(spec, "schema", "OneOfBar"),
+            ],
+            "discriminator": {
+                "propertyName": "type",
+                "mapping": {
+                    "foo": build_ref(spec, "schema", "OneOfFoo")["$ref"],
+                    "bar": build_ref(spec, "schema", "OneOfBar")["$ref"],
+                },
+            },
+        }
+        validate_spec(spec)
+
+    @pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
+    def test_spec_nested_oneof_schema(self, openapi_version):
+        spec = APISpec(
+            title="OneOf Nested Test",
+            version="1.0.0",
+            openapi_version=openapi_version,
+            plugins=[MarshmallowPlugin()],
+        )
+        spec.components.schema("Parent", schema=ExampleParentSchema)
+
+        schemas = get_schemas(spec)
+        assert "Parent" in schemas
+        assert "ExampleOneOf" in schemas
+        assert "OneOfFoo" in schemas
+        assert "OneOfBar" in schemas
+        assert schemas["Parent"]["properties"]["item"] == build_ref(
+            spec, "schema", "ExampleOneOf"
+        )
+        validate_spec(spec)
+
+    @pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
+    def test_spec_operations_with_oneof(self, openapi_version):
+        spec = APISpec(
+            title="OneOf Operations Test",
+            version="1.0.0",
+            openapi_version=openapi_version,
+            plugins=[MarshmallowPlugin()],
+        )
+        spec.path(
+            "/items",
+            operations={
+                "post": {
+                    "requestBody": {
+                        "content": {"application/json": {"schema": ExampleOneOf}}
+                    },
+                    "responses": {
+                        "200": {
+                            "description": "success",
+                            "content": {"application/json": {"schema": ExampleOneOf}},
+                        }
+                    },
+                }
+            },
+        )
+        schemas = get_schemas(spec)
+        assert "ExampleOneOf" in schemas
+        assert "OneOfFoo" in schemas
+        assert "OneOfBar" in schemas
+        paths = get_paths(spec)
+        assert paths["/items"]["post"]["requestBody"]["content"]["application/json"][
+            "schema"
+        ] == build_ref(spec, "schema", "ExampleOneOf")
+        validate_spec(spec)
+
+    @pytest.mark.parametrize("openapi_version", ("3.0.2", "3.1.0"))
+    def test_spec_anonymous_oneof_inline(self, openapi_version):
+        spec = APISpec(
+            title="Anonymous OneOf",
+            version="1.0.0",
+            openapi_version=openapi_version,
+            plugins=[MarshmallowPlugin(schema_name_resolver=lambda s: None)],
+        )
+        spec.path(
+            "/items",
+            operations={
+                "post": {
+                    "requestBody": {
+                        "content": {"application/json": {"schema": ExampleOneOf}}
+                    },
+                    "responses": {"200": {"description": "success"}},
+                }
+            },
+        )
+        paths = get_paths(spec)
+        schema = paths["/items"]["post"]["requestBody"]["content"]["application/json"][
+            "schema"
+        ]
+        assert "oneOf" in schema
+        assert len(schema["oneOf"]) == 2
+        assert schema["discriminator"] == {"propertyName": "type"}
+        validate_spec(spec)

tests/test_ext_marshmallow_openapi.py

@@ -717,3 +717,111 @@ def test_properties(self, field, properties, spec):
         for attr, expected_value in properties.items():
             assert attr in result
             assert result[attr] == expected_value
+
+
+try:
+    from marshmallow_oneofschema import OneOfSchema
+
+    ONEOFSCHEMA_AVAILABLE = True
+except ImportError:
+    ONEOFSCHEMA_AVAILABLE = False
+
+
+if ONEOFSCHEMA_AVAILABLE:
+
+    class FooSchema(Schema):
+        foo = fields.Str()
+
+    class BarSchema(Schema):
+        bar = fields.Int()
+
+    class MyOneOf(OneOfSchema):
+        type_field = "type"
+        type_schemas = {
+            "foo": FooSchema,
+            "bar": BarSchema,
+        }
+
+    class CustomFieldOneOf(OneOfSchema):
+        type_field = "custom_type"
+        type_schemas = {
+            "foo": FooSchema,
+            "bar": BarSchema,
+        }
+
+    class TitledOneOf(OneOfSchema):
+        type_field = "type"
+        type_schemas = {
+            "foo": FooSchema,
+            "bar": BarSchema,
+        }
+
+        class Meta:
+            title = "PolymorphicOffer"
+            description = "A polymorphic offer"
+
+    class EmptyOneOf(OneOfSchema):
+        type_schemas: dict[str, type[Schema]] = {}
+
+
+@pytest.mark.skipif(
+    not ONEOFSCHEMA_AVAILABLE, reason="marshmallow-oneofschema not installed"
+)
+class TestOneOfSchemaToOpenAPI:
+    @pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
+    def test_oneof_schema_to_jsonschema(self, openapi):
+        res = openapi.schema2jsonschema(MyOneOf)
+        assert res == {
+            "oneOf": [
+                {"$ref": "#/components/schemas/Foo"},
+                {"$ref": "#/components/schemas/Bar"},
+            ],
+            "discriminator": {
+                "propertyName": "type",
+                "mapping": {
+                    "foo": "#/components/schemas/Foo",
+                    "bar": "#/components/schemas/Bar",
+                },
+            },
+        }
+
+    @pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
+    def test_oneof_schema_instance_to_jsonschema(self, openapi):
+        res = openapi.schema2jsonschema(MyOneOf())
+        assert res == {
+            "oneOf": [
+                {"$ref": "#/components/schemas/Foo"},
+                {"$ref": "#/components/schemas/Bar"},
+            ],
+            "discriminator": {
+                "propertyName": "type",
+                "mapping": {
+                    "foo": "#/components/schemas/Foo",
+                    "bar": "#/components/schemas/Bar",
+                },
+            },
+        }
+
+    @pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
+    def test_oneof_custom_type_field(self, openapi):
+        res = openapi.schema2jsonschema(CustomFieldOneOf)
+        assert res["discriminator"]["propertyName"] == "custom_type"
+
+    @pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
+    def test_oneof_meta_title_and_description(self, openapi):
+        res = openapi.schema2jsonschema(TitledOneOf)
+        assert res["title"] == "PolymorphicOffer"
+        assert res["description"] == "A polymorphic offer"
+
+    @pytest.mark.parametrize("openapi", ("3.0.0", "3.1.0"), indirect=True)
+    def test_oneof_empty_type_schemas(self, openapi):
+        res = openapi.schema2jsonschema(EmptyOneOf)
+        assert res == {"type": "object", "properties": {}}
+
+    @pytest.mark.parametrize("openapi", ("2.0",), indirect=True)
+    def test_oneof_openapi_v2_raises_error(self, openapi):
+        with pytest.raises(
+            exceptions.APISpecError,
+            match="OneOfSchema is only supported in OpenAPI 3.0.0 and above",
+        ):
+            openapi.schema2jsonschema(MyOneOf)