Skip to content

Commit 9b40ee7

Browse files
Fix enum member names for object values to use title/name/const keys (#2596)
* Add support for enum members with object values in JSON schema * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 39a966f commit 9b40ee7

4 files changed

Lines changed: 103 additions & 0 deletions

File tree

src/datamodel_code_generator/parser/jsonschema.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1539,6 +1539,17 @@ def parse_enum_as_literal(self, obj: JsonSchemaObject) -> DataType:
15391539
"""Parse enum values as a Literal type."""
15401540
return self.data_type(literals=[i for i in obj.enum if i is not None])
15411541

1542+
@classmethod
1543+
def _get_field_name_from_dict_enum(cls, enum_part: dict[str, Any], index: int) -> str:
1544+
"""Extract field name from dict enum value using title, name, or const keys."""
1545+
if enum_part.get("title"):
1546+
return str(enum_part["title"])
1547+
if enum_part.get("name"):
1548+
return str(enum_part["name"])
1549+
if "const" in enum_part:
1550+
return str(enum_part["const"])
1551+
return f"value_{index}"
1552+
15421553
def parse_enum(
15431554
self,
15441555
name: str,
@@ -1575,6 +1586,8 @@ def parse_enum(
15751586
default = enum_part
15761587
if obj.x_enum_varnames:
15771588
field_name = obj.x_enum_varnames[i]
1589+
elif isinstance(enum_part, dict):
1590+
field_name = self._get_field_name_from_dict_enum(enum_part, i)
15781591
else:
15791592
prefix = obj.type if isinstance(obj.type, str) else type(enum_part).__name__
15801593
field_name = f"{prefix}_{enum_part}"
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# generated by datamodel-codegen:
2+
# filename: enum_object_values.json
3+
# timestamp: 2019-07-26T00:00:00+00:00
4+
5+
from __future__ import annotations
6+
7+
from enum import Enum
8+
from typing import Optional
9+
10+
from pydantic import BaseModel
11+
12+
13+
class CountryWithTitle(Enum):
14+
USA = {'title': 'USA', 'code': 'US'}
15+
CANADA = {'title': 'CANADA', 'code': 'CA'}
16+
17+
18+
class CountryWithName(Enum):
19+
USA = {'name': 'USA', 'code': 'US'}
20+
CANADA = {'name': 'CANADA', 'code': 'CA'}
21+
22+
23+
class CountryWithConst(Enum):
24+
us_value = {'const': 'us_value'}
25+
ca_value = {'const': 'ca_value'}
26+
27+
28+
class CountryNoIdentifier(Enum):
29+
value_0 = {'code': 'US'}
30+
value_1 = {'code': 'CA'}
31+
32+
33+
class CountryDuplicateTitle(Enum):
34+
SAME = {'title': 'SAME', 'code': 'US'}
35+
SAME_1 = {'title': 'SAME', 'code': 'CA'}
36+
37+
38+
class EnumObjectTest(BaseModel):
39+
country_with_title: Optional[CountryWithTitle] = None
40+
country_with_name: Optional[CountryWithName] = None
41+
country_with_const: Optional[CountryWithConst] = None
42+
country_no_identifier: Optional[CountryNoIdentifier] = None
43+
country_duplicate_title: Optional[CountryDuplicateTitle] = None
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
{
2+
"$schema": "http://json-schema.org/draft-07/schema#",
3+
"title": "EnumObjectTest",
4+
"type": "object",
5+
"properties": {
6+
"country_with_title": {
7+
"enum": [
8+
{"title": "USA", "code": "US"},
9+
{"title": "CANADA", "code": "CA"}
10+
]
11+
},
12+
"country_with_name": {
13+
"enum": [
14+
{"name": "USA", "code": "US"},
15+
{"name": "CANADA", "code": "CA"}
16+
]
17+
},
18+
"country_with_const": {
19+
"enum": [
20+
{"const": "us_value"},
21+
{"const": "ca_value"}
22+
]
23+
},
24+
"country_no_identifier": {
25+
"enum": [
26+
{"code": "US"},
27+
{"code": "CA"}
28+
]
29+
},
30+
"country_duplicate_title": {
31+
"enum": [
32+
{"title": "SAME", "code": "US"},
33+
{"title": "SAME", "code": "CA"}
34+
]
35+
}
36+
}
37+
}

tests/main/jsonschema/test_main_jsonschema.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2977,6 +2977,16 @@ def test_main_jsonschema_empty_items_array(output_file: Path) -> None:
29772977
)
29782978

29792979

2980+
def test_main_jsonschema_enum_object_values(output_file: Path) -> None:
2981+
"""Test that enum with object values uses title/name/const for member names (issue #1620)."""
2982+
run_main_and_assert(
2983+
input_path=JSON_SCHEMA_DATA_PATH / "enum_object_values.json",
2984+
output_path=output_file,
2985+
input_file_type="jsonschema",
2986+
assert_func=assert_file_content,
2987+
)
2988+
2989+
29802990
def test_main_jsonschema_collapse_root_models_empty_union(output_file: Path) -> None:
29812991
"""Test that collapse-root-models with empty union fallback generates Any instead of invalid Union syntax.
29822992

0 commit comments

Comments
 (0)