Skip to content

Commit 54c3ed9

Browse files
siminn-arnorgjsiminn-arnorgj
authored
fix: Handle class name prefix correctly in GraphQL parser (#2926)
* initial idea * add test for unions * add test for scalars * rename file
1 parent 58e73ed commit 54c3ed9

6 files changed

Lines changed: 250 additions & 4 deletions

File tree

‎src/datamodel_code_generator/model/scalar.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,9 @@ def __init__( # noqa: PLR0913
6767
extra_template_data[scalar_name] = defaultdict(dict)
6868

6969
# py_type
70-
py_type = extra_template_data[scalar_name].get(
70+
py_type = extra_template_data[reference.original_name].get(
7171
"py_type",
72-
DEFAULT_GRAPHQL_SCALAR_TYPES.get(reference.name, DEFAULT_GRAPHQL_SCALAR_TYPE),
72+
DEFAULT_GRAPHQL_SCALAR_TYPES.get(reference.original_name, DEFAULT_GRAPHQL_SCALAR_TYPE),
7373
)
7474
extra_template_data[scalar_name]["py_type"] = py_type
7575

‎src/datamodel_code_generator/parser/graphql.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -479,8 +479,10 @@ def parse_input_object(self, input_graphql_object: graphql.GraphQLInputObjectTyp
479479

480480
def parse_union(self, union_object: graphql.GraphQLUnionType) -> None:
481481
"""Parse a GraphQL union type and add it to results."""
482-
fields = [self.data_model_field_type(name=type_.name, data_type=DataType()) for type_ in union_object.types]
483-
482+
fields = [
483+
self.data_model_field_type(name=self.references[type_.name].name, data_type=DataType())
484+
for type_ in union_object.types
485+
]
484486
data_model_type = self.data_model_union_type(
485487
reference=self.references[union_object.name],
486488
fields=fields,

‎tests/data/expected/main/graphql/simple_star_wars_class_name_prefix.py‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
# generated by datamodel-codegen:
2+
# filename: simple-star-wars.graphql
3+
# timestamp: 2019-07-26T00:00:00+00:00
4+
5+
from __future__ import annotations
6+
7+
from typing import Literal, TypeAlias
8+
9+
from pydantic import BaseModel, Field
10+
11+
FooBoolean: TypeAlias = bool
12+
"""
13+
The `Boolean` scalar type represents `true` or `false`.
14+
"""
15+
16+
17+
FooID: TypeAlias = str
18+
"""
19+
The `ID` scalar type represents a unique identifier, often used to refetch an object or as key for a cache. The ID type appears in a JSON response as a String; however, it is not intended to be human-readable. When expected as an input type, any string (such as `"4"`) or integer (such as `4`) input value will be accepted as an ID.
20+
"""
21+
22+
23+
FooInt: TypeAlias = int
24+
"""
25+
The `Int` scalar type represents non-fractional signed whole numeric values. Int can represent values between -(2^31) and 2^31 - 1.
26+
"""
27+
28+
29+
FooString: TypeAlias = str
30+
"""
31+
The `String` scalar type represents textual data, represented as UTF-8 character sequences. The String type is most often used by GraphQL to represent free-form human-readable text.
32+
"""
33+
34+
35+
class FooFilm(BaseModel):
36+
characters: list[FooPerson]
37+
characters_ids: list[FooID]
38+
director: FooString
39+
episode_id: FooInt
40+
id: FooID
41+
opening_crawl: FooString
42+
planets: list[FooPlanet]
43+
planets_ids: list[FooID]
44+
producer: FooString | None = None
45+
release_date: FooString
46+
species: list[FooSpecies]
47+
species_ids: list[FooID]
48+
starships: list[FooStarship]
49+
starships_ids: list[FooID]
50+
title: FooString
51+
vehicles: list[FooVehicle]
52+
vehicles_ids: list[FooID]
53+
typename__: Literal['Film'] | None = Field('Film', alias='__typename')
54+
55+
56+
class FooPerson(BaseModel):
57+
birth_year: FooString | None = None
58+
eye_color: FooString | None = None
59+
films: list[FooFilm]
60+
films_ids: list[FooID]
61+
gender: FooString | None = None
62+
hair_color: FooString | None = None
63+
height: FooInt | None = None
64+
homeworld: FooPlanet | None = None
65+
homeworld_id: FooID | None = None
66+
id: FooID
67+
mass: FooInt | None = None
68+
name: FooString
69+
skin_color: FooString | None = None
70+
species: list[FooSpecies]
71+
species_ids: list[FooID]
72+
starships: list[FooStarship]
73+
starships_ids: list[FooID]
74+
vehicles: list[FooVehicle]
75+
vehicles_ids: list[FooID]
76+
typename__: Literal['Person'] | None = Field('Person', alias='__typename')
77+
78+
79+
class FooPlanet(BaseModel):
80+
climate: FooString | None = None
81+
diameter: FooString | None = None
82+
films: list[FooFilm]
83+
films_ids: list[FooID]
84+
gravity: FooString | None = None
85+
id: FooID
86+
name: FooString
87+
orbital_period: FooString | None = None
88+
population: FooString | None = None
89+
residents: list[FooPerson]
90+
residents_ids: list[FooID]
91+
rotation_period: FooString | None = None
92+
surface_water: FooString | None = None
93+
terrain: FooString | None = None
94+
typename__: Literal['Planet'] | None = Field('Planet', alias='__typename')
95+
96+
97+
class FooSpecies(BaseModel):
98+
average_height: FooString | None = None
99+
average_lifespan: FooString | None = None
100+
classification: FooString | None = None
101+
designation: FooString | None = None
102+
eye_colors: FooString | None = None
103+
films: list[FooFilm]
104+
films_ids: list[FooID]
105+
hair_colors: FooString | None = None
106+
id: FooID
107+
language: FooString | None = None
108+
name: FooString
109+
people: list[FooPerson]
110+
people_ids: list[FooID]
111+
skin_colors: FooString | None = None
112+
typename__: Literal['Species'] | None = Field('Species', alias='__typename')
113+
114+
115+
class FooStarship(BaseModel):
116+
MGLT: FooString | None = None
117+
cargo_capacity: FooString | None = None
118+
consumables: FooString | None = None
119+
cost_in_credits: FooString | None = None
120+
crew: FooString | None = None
121+
films: list[FooFilm]
122+
films_ids: list[FooID]
123+
hyperdrive_rating: FooString | None = None
124+
id: FooID
125+
length: FooString | None = None
126+
manufacturer: FooString | None = None
127+
max_atmosphering_speed: FooString | None = None
128+
model: FooString | None = None
129+
name: FooString
130+
passengers: FooString | None = None
131+
pilots: list[FooPerson]
132+
pilots_ids: list[FooID]
133+
starship_class: FooString | None = None
134+
typename__: Literal['Starship'] | None = Field('Starship', alias='__typename')
135+
136+
137+
class FooVehicle(BaseModel):
138+
cargo_capacity: FooString | None = None
139+
consumables: FooString | None = None
140+
cost_in_credits: FooString | None = None
141+
crew: FooString | None = None
142+
films: list[FooFilm]
143+
films_ids: list[FooID]
144+
id: FooID
145+
length: FooString | None = None
146+
manufacturer: FooString | None = None
147+
max_atmosphering_speed: FooString | None = None
148+
model: FooString | None = None
149+
name: FooString
150+
passengers: FooString | None = None
151+
pilots: list[FooPerson]
152+
pilots_ids: list[FooID]
153+
vehicle_class: FooString | None = None
154+
typename__: Literal['Vehicle'] | None = Field('Vehicle', alias='__typename')
155+
156+
157+
FooFilm.update_forward_refs()
158+
FooPerson.update_forward_refs()

‎tests/data/expected/parser/graphql/union_with_prefix.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# generated by datamodel-codegen:
2+
# filename: union.graphql
3+
# timestamp: 2019-07-26T00:00:00+00:00
4+
5+
from __future__ import annotations
6+
7+
from typing import Literal, TypeAlias, Union
8+
9+
from pydantic import BaseModel, Field
10+
11+
FooBoolean: TypeAlias = bool
12+
"""
13+
The `Boolean` scalar type represents `true` or `false`.
14+
"""
15+
16+
17+
FooID: TypeAlias = str
18+
"""
19+
The `ID` scalar type represents a unique identifier, often used to refetch an object or as key for a cache. The ID type appears in a JSON response as a String; however, it is not intended to be human-readable. When expected as an input type, any string (such as `"4"`) or integer (such as `4`) input value will be accepted as an ID.
20+
"""
21+
22+
23+
FooInt: TypeAlias = int
24+
"""
25+
The `Int` scalar type represents non-fractional signed whole numeric values. Int can represent values between -(2^31) and 2^31 - 1.
26+
"""
27+
28+
29+
FooString: TypeAlias = str
30+
"""
31+
The `String` scalar type represents textual data, represented as UTF-8 character sequences. The String type is most often used by GraphQL to represent free-form human-readable text.
32+
"""
33+
34+
35+
class FooIResource(BaseModel):
36+
id: FooID
37+
typename__: Literal['IResource'] | None = Field('IResource', alias='__typename')
38+
39+
40+
class FooCar(FooIResource):
41+
id: FooID
42+
passengerCapacity: FooInt
43+
typename__: Literal['Car'] | None = Field('Car', alias='__typename')
44+
45+
46+
class FooEmployee(FooIResource):
47+
firstName: FooString | None = None
48+
id: FooID
49+
lastName: FooString | None = None
50+
typename__: Literal['Employee'] | None = Field('Employee', alias='__typename')
51+
52+
53+
FooResource: TypeAlias = Union[
54+
'FooCar',
55+
'FooEmployee',
56+
]
57+
58+
59+
FooTechnicalResource: TypeAlias = FooCar

‎tests/main/graphql/test_main_graphql.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,6 +762,21 @@ def test_main_graphql_dataclass_frozen_keyword_only(output_file: Path) -> None:
762762
)
763763

764764

765+
def test_main_graphql_class_name_prefix(output_file: Path) -> None:
766+
"""Test GraphQL code generation with class name prefixing."""
767+
run_main_and_assert(
768+
input_path=GRAPHQL_DATA_PATH / "simple-star-wars.graphql",
769+
output_path=output_file,
770+
input_file_type="graphql",
771+
assert_func=assert_file_content,
772+
expected_file="simple_star_wars_class_name_prefix.py",
773+
extra_args=[
774+
"--class-name-prefix",
775+
"Foo",
776+
],
777+
)
778+
779+
765780
def test_main_graphql_union_snake_case_field(output_file: Path) -> None:
766781
"""Test that union type references are not converted to snake_case."""
767782
run_main_and_assert(

‎tests/parser/test_graphql.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,18 @@ def test_graphql_union_commented(output_file: Path) -> None:
5555
)
5656

5757

58+
def test_graphql_union_with_prefix(output_file: Path) -> None:
59+
"""Test parsing GraphQL union with class name prefix (Unions should reference prefixed class names)."""
60+
run_main_and_assert(
61+
input_path=GRAPHQL_DATA_PATH / "union.graphql",
62+
output_path=output_file,
63+
input_file_type="graphql",
64+
assert_func=assert_file_content,
65+
expected_file="union_with_prefix.py",
66+
extra_args=["--class-name-prefix", "Foo"],
67+
)
68+
69+
5870
@pytest.mark.parametrize(
5971
("frozen_dataclasses", "keyword_only", "parser_dataclass_args", "kwargs_dataclass_args", "expected"),
6072
[

0 commit comments

Comments
 (0)