Skip to content

Commit

Permalink
Fix default for annotated field in pydantic v2 (#1498)
Browse files Browse the repository at this point in the history
* Fix default for annotated field in pydantic v2

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move v2 logic to v2 model

---------

Co-authored-by: ferris <ferris@devdroplets.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Koudai Aono <koxudaxi@gmail.com>
  • Loading branch information
4 people committed Nov 25, 2023
1 parent 5ccc441 commit a46fe94
Show file tree
Hide file tree
Showing 5 changed files with 127 additions and 8 deletions.
7 changes: 6 additions & 1 deletion datamodel_code_generator/model/pydantic/base_model.py
Expand Up @@ -125,6 +125,11 @@ def _process_data_in_str(self, data: Dict[str, Any]) -> None:
if self.const:
data['const'] = True

def _process_annotated_field_arguments(
self, field_arguments: List[str]
) -> List[str]:
return field_arguments

def __str__(self) -> str:
data: Dict[str, Any] = {
k: v for k, v in self.extras.items() if k not in self._EXCLUDE_FIELD_KEYS
Expand Down Expand Up @@ -180,7 +185,7 @@ def __str__(self) -> str:
return ''

if self.use_annotated:
pass
field_arguments = self._process_annotated_field_arguments(field_arguments)
elif self.required:
field_arguments = ['...', *field_arguments]
elif default_factory:
Expand Down
14 changes: 14 additions & 0 deletions datamodel_code_generator/model/pydantic_v2/base_model.py
Expand Up @@ -100,6 +100,20 @@ def _process_data_in_str(self, data: Dict[str, Any]) -> None:
# unique_items is not supported in pydantic 2.0
data.pop('unique_items', None)

def _process_annotated_field_arguments(
self, field_arguments: List[str]
) -> List[str]:
if not self.required:
if self.use_default_kwarg:
return [
f'default={repr(self.default)}',
*field_arguments,
]
else:
# TODO: Allow '=' style default for v1?
return [f'{repr(self.default)}', *field_arguments]
return field_arguments


class ConfigAttribute(NamedTuple):
from_: str
Expand Down
Expand Up @@ -24,7 +24,7 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme
{%- else %}
{{ field.name }}: {{ field.type_hint }}
{%- endif %}
{%- if not field.required or field.data_type.is_optional or field.nullable
{%- if (not field.required or field.data_type.is_optional or field.nullable) and not field.annotated
%} = {{ field.represented_default }}
{%- endif -%}
{%- endif %}
Expand Down
@@ -0,0 +1,92 @@
# generated by datamodel-codegen:
# filename: api_constrained.yaml
# timestamp: 2019-07-26T00:00:00+00:00

from __future__ import annotations

from typing import Annotated, List, Optional, Union

from pydantic import AnyUrl, BaseModel, Field, RootModel


class Pet(BaseModel):
id: Annotated[int, Field(ge=0, le=9223372036854775807)]
name: Annotated[str, Field(max_length=256)]
tag: Annotated[Optional[str], Field(None, max_length=64)]


class Pets(RootModel[List[Pet]]):
root: Annotated[List[Pet], Field(max_length=10, min_length=1)]


class UID(RootModel[int]):
root: Annotated[int, Field(ge=0)]


class Phone(RootModel[str]):
root: Annotated[str, Field(min_length=3)]


class FaxItem(RootModel[str]):
root: Annotated[str, Field(min_length=3)]


class User(BaseModel):
id: Annotated[int, Field(ge=0)]
name: Annotated[str, Field(max_length=256)]
tag: Annotated[Optional[str], Field(None, max_length=64)]
uid: UID
phones: Annotated[Optional[List[Phone]], Field(None, max_length=10)]
fax: Optional[List[FaxItem]] = None
height: Annotated[Optional[Union[int, float]], Field(None, ge=1.0, le=300.0)]
weight: Annotated[Optional[Union[float, int]], Field(None, ge=1.0, le=1000.0)]
age: Annotated[Optional[int], Field(None, gt=0, le=200)]
rating: Annotated[Optional[float], Field(None, gt=0.0, le=5.0)]


class Users(RootModel[List[User]]):
root: List[User]


class Id(RootModel[str]):
root: str


class Rules(RootModel[List[str]]):
root: List[str]


class Error(BaseModel):
code: int
message: str


class Api(BaseModel):
apiKey: Annotated[
Optional[str],
Field(None, description='To be used as a dataset parameter value'),
]
apiVersionNumber: Annotated[
Optional[str],
Field(None, description='To be used as a version parameter value'),
]
apiUrl: Annotated[
Optional[AnyUrl],
Field(None, description="The URL describing the dataset's fields"),
]
apiDocumentationUrl: Annotated[
Optional[AnyUrl],
Field(None, description='A URL to the API console for each API'),
]


class Apis(RootModel[List[Api]]):
root: List[Api]


class Event(BaseModel):
name: Optional[str] = None


class Result(BaseModel):
event: Optional[Event] = None
20 changes: 14 additions & 6 deletions tests/test_main.py
Expand Up @@ -4251,7 +4251,17 @@ def test_jsonschema_without_titles_use_title_as_name():


@freeze_time('2019-07-26')
def test_main_use_annotated_with_field_constraints():
@pytest.mark.parametrize(
'output_model,expected_output',
[
('pydantic.BaseModel', 'main_use_annotated_with_field_constraints'),
(
'pydantic_v2.BaseModel',
'main_use_annotated_with_field_constraints_pydantic_v2',
),
],
)
def test_main_use_annotated_with_field_constraints(output_model, expected_output):
with TemporaryDirectory() as output_dir:
output_file: Path = Path(output_dir) / 'output.py'
return_code: Exit = main(
Expand All @@ -4264,16 +4274,14 @@ def test_main_use_annotated_with_field_constraints():
'--use-annotated',
'--target-python-version',
'3.9',
'--output-model',
output_model,
]
)
assert return_code == Exit.OK
assert (
output_file.read_text()
== (
EXPECTED_MAIN_PATH
/ 'main_use_annotated_with_field_constraints'
/ 'output.py'
).read_text()
== (EXPECTED_MAIN_PATH / expected_output / 'output.py').read_text()
)


Expand Down

0 comments on commit a46fe94

Please sign in to comment.