Skip to content

Commit 0db5cdc

Browse files
authored
Fix msgspec enum discriminator tag (#3170)
1 parent 4d0d9a9 commit 0db5cdc

4 files changed

Lines changed: 74 additions & 1 deletion

File tree

src/datamodel_code_generator/parser/base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1794,7 +1794,7 @@ def get_enum_from_base(
17941794
has_one_literal = True
17951795
if isinstance(discriminator_model, msgspec_model.Struct): # pragma: no cover
17961796
discriminator_model.add_base_class_kwarg("tag_field", f"'{field_name}'")
1797-
discriminator_model.add_base_class_kwarg("tag", discriminator_field.represented_default)
1797+
discriminator_model.add_base_class_kwarg("tag", repr(expected_value))
17981798
discriminator_field.extras["is_classvar"] = True
17991799
# Found the discriminator field, no need to keep looking
18001800
break
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# generated by datamodel-codegen:
2+
# filename: discriminator_enum_single_value_msgspec.yaml
3+
# timestamp: 2019-07-26T00:00:00+00:00
4+
5+
from __future__ import annotations
6+
7+
from typing import TypeAlias
8+
9+
from msgspec import Struct
10+
11+
12+
class StartedEvent(Struct, tag_field='type', tag='started'):
13+
pass
14+
15+
16+
class StoppedEvent(Struct, tag_field='type', tag='stopped'):
17+
pass
18+
19+
20+
Event: TypeAlias = StartedEvent | StoppedEvent
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
openapi: 3.1.0
2+
info:
3+
title: Repro
4+
version: 0.0.1
5+
paths: {}
6+
components:
7+
schemas:
8+
Event:
9+
discriminator:
10+
propertyName: type
11+
oneOf:
12+
- $ref: '#/components/schemas/StartedEvent'
13+
- $ref: '#/components/schemas/StoppedEvent'
14+
StartedEvent:
15+
type: object
16+
properties:
17+
type:
18+
type: string
19+
enum:
20+
- started
21+
required:
22+
- type
23+
StoppedEvent:
24+
type: object
25+
properties:
26+
type:
27+
type: string
28+
enum:
29+
- stopped
30+
required:
31+
- type

tests/main/openapi/test_main_openapi.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,28 @@ def test_main_openapi_discriminator_enum_single_value_anyof_use_enum(output_file
331331
)
332332

333333

334+
def test_main_openapi_discriminator_enum_single_value_msgspec(output_file: Path) -> None:
335+
"""Single-value enum discriminator is used as the msgspec tag."""
336+
run_main_and_assert(
337+
input_path=OPEN_API_DATA_PATH / "discriminator_enum_single_value_msgspec.yaml",
338+
output_path=output_file,
339+
input_file_type="openapi",
340+
assert_func=assert_file_content,
341+
expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "enum_single_value_msgspec.py",
342+
extra_args=[
343+
"--target-python-version",
344+
"3.10",
345+
"--output-model-type",
346+
"msgspec.Struct",
347+
"--enum-field-as-literal",
348+
"all",
349+
"--use-one-literal-as-default",
350+
"--disable-warnings",
351+
],
352+
force_exec_validation=True,
353+
)
354+
355+
334356
def test_main_openapi_discriminator_with_properties(output_file: Path) -> None:
335357
"""Test OpenAPI generation with discriminator properties."""
336358
run_main_and_assert(

0 commit comments

Comments
 (0)