Incorrect handling of Union types

[Amwam]

It seems that the API spec is not properly generated when using Union typehints.
The below snippet is on python 3.10 and 3.11

from typing import Literal

from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
from marshmallow_dataclass import dataclass


class BaseItemMixin:
    name: str


@dataclass
class BaseItem(BaseItemMixin):
    pass


@dataclass
class Food(BaseItem):
    label: str
    name: Literal["food"] = "food"


@dataclass
class Drink(BaseItem):
    size: int
    name: Literal["drink"] = "drink"


@dataclass
class MainSchema:
    id: int
    name: str
    items: list[Food | Drink]


spec = APISpec(
    title="API",
    version="1.0.0",
    openapi_version="3.0.2",
    plugins=[(MarshmallowPlugin())],
)


spec.path(
    "/endpoint",
    operations={
        "post": {
            "parameters": [
                {
                    "name": "Args",
                    "schema": MainSchema.Schema(),
                    "in": "body",
                }
            ],
            "responses": {
                "200": {
                    "content": {"application/json": {}},
                    "description": "OK",
                },
            },
        }
    },
)

if __name__ == "__main__":
    print(spec.to_yaml())

The outputted yaml contains the schema for Main as:

components:
  schemas:
    Main:
      type: object
      properties:
        name:
          type: string
        id:
          type: integer
        items:
          type: array
          items: {}
      required:
      - id
      - items
      - name

Items is correctly marked as an array, but the items is not correct. I'd expect oneOf being used here (although i'm not 100%).
The same happens using item: Food | Drink.

I believe the issue is within this library, and not marshmallow_dataclass, as i'm able to deserialise serialise with the schemas without issue

    data = MainSchema(
        id=1,
        name="test",
        items=[
            Food(name="food", label="test"),
            Drink(name="drink", size=1),
        ],
    )
    data_as_dict = MainSchema.Schema().dumps(data)
    loaded_data = MainSchema.Schema().loads(data_as_dict)
    print("original", data)
	print("deserialised", loaded_data)
    print("serialised", data_as_dict)
   
# original MainSchema(id=1, name='test', items=[Food(label='test', name='food'), Drink(size=1, name='drink')])
# deserialised MainSchema(id=1, name='test', items=[Food(label='test', name='food'), Drink(size=1, name='drink')])
# serialised {"id": 1, "items": [{"label": "test", "name": "food"}, {"size": 1, "name": "drink"}], "name": "test"}

Note: I've also tried without the inheritance involved, and the result is the same.

[Amwam]

I think i may have found a solution, by adding a custom attribute function

def convert_union_to_spec(self, field, **kwargs):
    ret = {}
    if isinstance(field, marshmallow_dataclass.union_field.Union):
        ret["oneOf"] = [self.field2property(f[1]) for f in field.union_fields]

    return ret


ma_plugin.converter.add_attribute_function(convert_union_to_spec)

Would still appreciate some guidance if this is the right approach. Thanks!

[lafrech]

I'm not familiar with dataclass so I can't provide an educated advice.

This custom attribute function looks good to me.

[vkruglik-aka]

Hi @Amwam, thanks for sharing your example with ma_plugin.converter.add_attribute_function(convert_union_to_spec), but this ticket doesn't show what ma_plugin is and how it's created, so I am scratching my head, trying to figure it out.

[Amwam]

@vkruglik-aka Hi, the ma_plugin is an instance of MarshmallowPlugin, the code is along the lines of:

from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin

ma_plugin = MarshmallowPlugin()

def convert_union_to_spec(self, field, **kwargs):
    ret = {}
    if isinstance(field, marshmallow_dataclass.union_field.Union):
        ret["oneOf"] = [self.field2property(f[1]) for f in field.union_fields]

    return ret


ma_plugin.converter.add_attribute_function(convert_union_to_spec)

spec = APISpec(
    title="My API",
    version="1.0.0",
    openapi_version="3.0.2",
    plugins=[ma_plugin],
)

[vkruglik-aka]

Thank you @Amwam, thank you you for following up. Please see code correction below. I figured this out once I learned more about marshmallow API.

Correction:

ma_plugin.converter.add_attribute_function(convert_union_to_spec) needs to be called after spec = APISpec(...). If called before, then ma_plugin.converter is still None, so the call to ma_plugin.converter.add_attribute_function() fails.