Repository: koxudaxi/datamodel-code-generator

Issue:
Wrong parent class for pydantic V2 root models 
**Describe the bug**
Generator uses `pydantic.BaseModel` as parent class for root model instead of `pydantic.RootModel`

Example schema (`custom_id.yaml`):
```yaml
openapi: 3.0.0
components:
  schemas:
    CustomId:
      description: My custom ID
      type: string
      format: uuid
```
Used commandline:
```
$ datamodel-codegen --input custom_id.yaml --output-model-type pydantic_v2.BaseModel --output model.py
```
Contents of `model.py`:
```python
from __future__ import annotations

from uuid import UUID

from pydantic import BaseModel, Field


class CustomId(BaseModel):
    root: UUID = Field(..., description='My custom ID')
```

**Expected behavior**
```python
from __future__ import annotations

from uuid import UUID

from pydantic import RootModel, Field


class CustomId(RootModel):
    root: UUID = Field(..., description='My custom ID')
```

**Version:**
 - OS: [e.g. iOS]
 - Python version: 3.10.8
 - datamodel-code-generator version: 0.21.1

Wrong parent class for pydantic V2 root models 
**Describe the bug**
Generator uses `pydantic.BaseModel` as parent class for root model instead of `pydantic.RootModel`

Example schema (`custom_id.yaml`):
```yaml
openapi: 3.0.0
components:
  schemas:
    CustomId:
      description: My custom ID
      type: string
      format: uuid
```
Used commandline:
```
$ datamodel-codegen --input custom_id.yaml --output-model-type pydantic_v2.BaseModel --output model.py
```
Contents of `model.py`:
```python
from __future__ import annotations

from uuid import UUID

from pydantic import BaseModel, Field


class CustomId(BaseModel):
    root: UUID = Field(..., description='My custom ID')
```

**Expected behavior**
```python
from __future__ import annotations

from uuid import UUID

from pydantic import RootModel, Field


class CustomId(RootModel):
    root: UUID = Field(..., description='My custom ID')
```

**Version:**
 - OS: [e.g. iOS]
 - Python version: 3.10.8
 - datamodel-code-generator version: 0.21.1



Discussion:
@andreyLetenkovWefox 
Thank you for creating the issue.
I have tested your schema. But I can't reproduce the issue.
Do you use a custom template for root model?
> @andreyLetenkovWefox Thank you for creating the issue. I have tested your schema. But I can't reproduce the issue. Do you use a custom template for root model?

Sorry, I thought the problem could be reproduced via a CLI generator call, but it really only appears when using the python package:
```python
from pathlib import Path

from datamodel_code_generator import DataModelType, generate

generate(
    input_=Path("custom_id.yaml"),
    output=Path("model.py"),
    output_model_type=DataModelType.PydanticV2BaseModel,
)
```

@andreyLetenkovWefox 
Thank you for creating the issue.
I have tested your schema. But I can't reproduce the issue.
Do you use a custom template for root model?
> @andreyLetenkovWefox Thank you for creating the issue. I have tested your schema. But I can't reproduce the issue. Do you use a custom template for root model?

Sorry, I thought the problem could be reproduced via a CLI generator call, but it really only appears when using the python package:
```python
from pathlib import Path

from datamodel_code_generator import DataModelType, generate

generate(
    input_=Path("custom_id.yaml"),
    output=Path("model.py"),
    output_model_type=DataModelType.PydanticV2BaseModel,
)
```


Resolving patch:
diff --git a/datamodel_code_generator/__main__.py b/datamodel_code_generator/__main__.py
--- a/datamodel_code_generator/__main__.py
+++ b/datamodel_code_generator/__main__.py
@@ -36,8 +36,10 @@
 import toml
 from pydantic import BaseModel
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 from datamodel_code_generator import (
-    DEFAULT_BASE_CLASS,
     DataModelType,
     Error,
     InputFileType,
@@ -470,9 +472,6 @@ def get(self, item: str) -> Any:
         def __getitem__(self, item: str) -> Any:
             return self.get(item)
 
-        def __setitem__(self, key: str, value: Any) -> None:
-            setattr(self, key, value)
-
         if TYPE_CHECKING:
 
             @classmethod
@@ -580,23 +579,21 @@ def validate_each_item(each_item: Any) -> Tuple[str, str]:
             return [validate_each_item(each_item) for each_item in value]
         return value  # pragma: no cover
 
-    @model_validator(mode='after')
-    def validate_root(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        values = cls._validate_use_annotated(values)
-        return cls._validate_base_class(values)
-
-    @classmethod
-    def _validate_use_annotated(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get('use_annotated'):
-            values['field_constraints'] = True
-        return values
+    if PYDANTIC_V2:
 
-    @classmethod
-    def _validate_base_class(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if 'base_class' not in values and 'output_model_type' in values:
-            if values['output_model_type'] != DataModelType.PydanticBaseModel.value:
-                values['base_class'] = ''
-        return values
+        @model_validator(mode='after')  # type: ignore
+        def validate_root(self: Self) -> Self:
+            if self.use_annotated:
+                self.field_constraints = True
+            return self
+
+    else:
+
+        @model_validator(mode='after')
+        def validate_root(cls, values: Any) -> Any:
+            if values.get('use_annotated'):
+                values['field_constraints'] = True
+            return values
 
     input: Optional[Union[Path, str]] = None
     input_file_type: InputFileType = InputFileType.Auto
@@ -605,7 +602,7 @@ def _validate_base_class(cls, values: Dict[str, Any]) -> Dict[str, Any]:
     debug: bool = False
     disable_warnings: bool = False
     target_python_version: PythonVersion = PythonVersion.PY_37
-    base_class: str = DEFAULT_BASE_CLASS
+    base_class: str = ''
     custom_template_dir: Optional[Path] = None
     extra_template_data: Optional[TextIOBase] = None
     validation: bool = False
@@ -666,9 +663,11 @@ def merge_args(self, args: Namespace) -> None:
             for f in self.get_fields()
             if getattr(args, f) is not None
         }
-        set_args = self._validate_use_annotated(set_args)
-        set_args = self._validate_base_class(set_args)
-        parsed_args = self.parse_obj(set_args)
+
+        if set_args.get('use_annotated'):
+            set_args['field_constraints'] = True
+
+        parsed_args = Config.parse_obj(set_args)
         for field_name in set_args:
             setattr(self, field_name, getattr(parsed_args, field_name))
 


Test patch:
diff --git a/tests/data/expected/main/main_openapi_custom_id_pydantic_v2/output.py b/tests/data/expected/main/main_openapi_custom_id_pydantic_v2/output.py
new file mode 100644
--- /dev/null
+++ b/tests/data/expected/main/main_openapi_custom_id_pydantic_v2/output.py
@@ -0,0 +1,18 @@
+# generated by datamodel-codegen:
+#   filename:  custom_id.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+from uuid import UUID
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class CustomId(RootModel):
+    root: UUID = Field(..., description='My custom ID')
+
+
+class Model(BaseModel):
+    custom_id: Optional[CustomId] = None
diff --git a/tests/data/expected/main/main_openapi_custom_id_pydantic_v2_custom_base/output.py b/tests/data/expected/main/main_openapi_custom_id_pydantic_v2_custom_base/output.py
new file mode 100644
--- /dev/null
+++ b/tests/data/expected/main/main_openapi_custom_id_pydantic_v2_custom_base/output.py
@@ -0,0 +1,20 @@
+# generated by datamodel-codegen:
+#   filename:  custom_id.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+from uuid import UUID
+
+from pydantic import Field, RootModel
+
+from custom_base import Base
+
+
+class CustomId(RootModel):
+    root: UUID = Field(..., description='My custom ID')
+
+
+class Model(Base):
+    custom_id: Optional[CustomId] = None
diff --git a/tests/data/openapi/custom_id.yaml b/tests/data/openapi/custom_id.yaml
new file mode 100644
--- /dev/null
+++ b/tests/data/openapi/custom_id.yaml
@@ -0,0 +1,12 @@
+openapi: 3.0.0
+components:
+  schemas:
+    CustomId:
+      description: My custom ID
+      type: string
+      format: uuid
+    Model:
+      type: object
+      properties:
+        custom_id:
+          $ref: "#/components/schemas/CustomId"
\ No newline at end of file
diff --git a/tests/test_main.py b/tests/test_main.py
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -5527,3 +5527,57 @@ def test_main_pydantic_v2():
             output_file.read_text()
             == (EXPECTED_MAIN_PATH / 'main_pydantic_v2' / 'output.py').read_text()
         )
+
+
+@freeze_time('2019-07-26')
+def test_main_openapi_custom_id_pydantic_v2():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(OPEN_API_DATA_PATH / 'custom_id.yaml'),
+                '--output',
+                str(output_file),
+                '--output-model-type',
+                'pydantic_v2.BaseModel',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (
+                EXPECTED_MAIN_PATH / 'main_openapi_custom_id_pydantic_v2' / 'output.py'
+            ).read_text()
+        )
+
+
+@pytest.mark.skipif(
+    not isort.__version__.startswith('4.'),
+    reason="isort 5.x don't sort pydantic modules",
+)
+@freeze_time('2019-07-26')
+def test_main_openapi_custom_id_pydantic_v2_custom_base():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(OPEN_API_DATA_PATH / 'custom_id.yaml'),
+                '--output',
+                str(output_file),
+                '--output-model-type',
+                'pydantic_v2.BaseModel',
+                '--base-class',
+                'custom_base.Base',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (
+                EXPECTED_MAIN_PATH
+                / 'main_openapi_custom_id_pydantic_v2_custom_base'
+                / 'output.py'
+            ).read_text()
+        )