1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
|
"""Tests for Pydantic v2 dataclass generation."""
from __future__ import annotations
from datamodel_code_generator.model import DataModelFieldBase
from datamodel_code_generator.model.pydantic_v2.dataclass import DataClass, DataModelField
from datamodel_code_generator.model.pydantic_v2.types import DataTypeManager
from datamodel_code_generator.reference import Reference
from datamodel_code_generator.types import DataType, Types
def test_data_class() -> None:
"""Test basic DataClass generation with required field."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)
data_class = DataClass(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
)
assert data_class.name == "test_model"
assert data_class.decorators == []
rendered = data_class.render()
assert "@dataclass" in rendered
assert "class test_model:" in rendered
assert "a: str" in rendered
def test_data_class_base_class() -> None:
"""Test DataClass generation with base class inheritance."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)
data_class = DataClass(
fields=[field],
base_classes=[Reference(name="Base", original_name="Base", path="Base")],
reference=Reference(name="test_model", path="test_model"),
)
assert data_class.name == "test_model"
assert data_class.decorators == []
rendered = data_class.render()
assert "@dataclass" in rendered
assert "class test_model(Base):" in rendered
assert "a: str" in rendered
def test_data_class_optional() -> None:
"""Test DataClass generation with field default value."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), default="'abc'", required=False)
data_class = DataClass(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
)
assert data_class.name == "test_model"
rendered = data_class.render()
assert "@dataclass" in rendered
assert "class test_model:" in rendered
# pydantic v2 uses Optional[str] for optional fields
assert "a:" in rendered
def test_data_class_get_data_type() -> None:
"""Test data type retrieval for DataClass fields."""
data_type_manager = DataTypeManager()
# Check that the type is correctly mapped
result = data_type_manager.get_data_type(Types.integer)
assert result.type == "int"
def test_data_class_field_ordering() -> None:
"""Test that fields are sorted: required first, defaults last."""
required_field = DataModelFieldBase(name="required_field", data_type=DataType(type="str"), required=True)
optional_field = DataModelFieldBase(
name="optional_field", data_type=DataType(type="str"), default="None", required=False
)
# Pass fields in wrong order (optional first)
data_class = DataClass(
fields=[optional_field, required_field],
reference=Reference(name="test_model", path="test_model"),
)
# Required field should come first after sorting
assert data_class.fields[0].name == "required_field"
assert data_class.fields[1].name == "optional_field"
def test_data_class_frozen() -> None:
"""Test DataClass with frozen=True."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)
data_class = DataClass(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
frozen=True,
)
assert data_class.dataclass_arguments.get("frozen") is True
assert "@dataclass(frozen=True)" in data_class.render()
def test_data_class_kw_only() -> None:
"""Test DataClass with keyword_only=True."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)
data_class = DataClass(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
keyword_only=True,
)
assert data_class.dataclass_arguments.get("kw_only") is True
assert "@dataclass(kw_only=True)" in data_class.render()
def test_data_class_with_description() -> None:
"""Test DataClass with docstring description."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)
data_class = DataClass(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
description="This is a test model.",
)
rendered = data_class.render()
assert "class test_model:" in rendered
assert '"""' in rendered
assert "This is a test model." in rendered
def test_data_model_field() -> None:
"""Test DataModelField creation."""
field = DataModelField(
name="test_field",
data_type=DataType(type="str"),
required=True,
)
assert field.name == "test_field"
assert field.required is True
def test_create_reuse_model() -> None:
"""Test creating a reuse model from existing DataClass."""
field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)
data_class = DataClass(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
frozen=True,
)
base_ref = Reference(name="BaseModel", path="base_model")
reuse_model = data_class.create_reuse_model(base_ref)
assert reuse_model.fields == []
assert len(reuse_model.base_classes) == 1
# base_classes are wrapped in BaseClassDataType which has a reference attribute
assert reuse_model.base_classes[0].reference is base_ref
assert reuse_model.frozen is True
|
