datamodel-code-generator icon indicating copy to clipboard operation
datamodel-code-generator copied to clipboard

Support for discriminated unions for the msgspec output

Open indrat opened this issue 2 years ago • 2 comments

Is your feature request related to a problem? Please describe. msgspec support for discriminated unions is not currently supported by datamodel-code-generator and it would be great to support discriminated unions.

Currently if generating discriminated union using the msgspec output, e.g from tests/data/jsonschema/discriminator_literals.json attempting to decode a yaml file with the generated code will result in at least two issues:

  • erroneousOptional type wrapping the type_ field; and
  • if the Optional type wrapper is removed then when attempting to use Response type to decode an input the following error is raised: TypeError: If a type union contains multiple Struct types, all Struct types must be tagged (via tagortag_fieldkwarg) - typetyping.Annotated[typing.Union[main.Type1, main.Type2], msgspec.Meta(title='Inner')] is not supported
$ datamodel-codegen --input tests/data/jsonschema/discriminator_literals.json --input-file-type jsonschema --output tests/data/expected/main/main_jsonschema_discriminator_literals_msgspec/output.py --disable-timestamp --enable-version-header --output-model-type msgspec.Struct --target-python 3.9

$ cat <<"EOF" >> tests/data/expected/main/main_jsonschema_discriminator_literals_msgspec/output.py 

if __name__ == "__main__":
    import msgspec
    d = """\
---
inner:
    type_: a
"""
    print(msgspec.yaml.decode(d, type=Response))
EOF

$ › python tests/data/expected/main/main_jsonschema_discriminator_literals_msgspec/output.py 
Traceback (most recent call last):
  File "tests/data/expected/main/main_jsonschema_discriminator_literals_msgspec/output.py ", line 31, in <module>
    print(msgspec.yaml.decode(d, type=Response))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/msgspec/yaml.py", line 161, in decode
    return _convert(
           ^^^^^^^^^
TypeError: If a type union contains multiple Struct types, all Struct types must be tagged (via `tag` or `tag_field` kwarg) - type `typing.Annotated[typing.Union[__main__.Type1, __main__.Type2], msgspec.Meta(title='Inner')]` is not supported

output.py with Optional type wrapper removed.

# generated by datamodel-codegen:
#   filename:  discriminator_literals.json
#   version:   0.0.0

from __future__ import annotations

from typing import Annotated, Literal, Union

from msgspec import Meta, Struct, field


class Type1(Struct):
    type_: Annotated[Literal['a'], Meta(title='Type ')] = 'a'


class Type2(Struct):
    type_: Annotated[Literal['b'], Meta(title='Type ')] = 'b'


class Response(Struct):
    inner: Annotated[Union[Type1, Type2], Meta(title='Inner')]


if __name__ == "__main__":
    import msgspec
    d = """\
---
inner:
    type_: a
"""
    print(msgspec.yaml.decode(d, type=Response))

Describe the solution you'd like To be able to define a discriminated union in jsonschema and generate the correct python/msgspec output. Preferably in the using the ClassVar format suggested by: https://github.com/jcrist/msgspec/issues/338#issuecomment-1453701495

Additional context Added example output test to: https://github.com/koxudaxi/datamodel-code-generator/commit/aaa3f7a598e6902c59f13b8ae0bf0adf4b92e37f unsure how to go about adding support for discriminators to the msgspec output.

indrat avatar Oct 07 '23 03:10 indrat

Is this complete with #1610 merged?

ncoghlan avatar Dec 03 '24 07:12 ncoghlan

I appear to also be running into this issue.

dhirschfeld avatar Mar 01 '25 10:03 dhirschfeld