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



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,
)
```

We hit this issue in a slightly different setting: if `output_data_model = "pydantic_v2.BaseModel"` is read from `pyproject.toml` and not passed on the command line, root models have the wrong base class. However, if `--output_data_model "pydantic_v2.BaseModel"` is passed on the command line, the root model base class is correct.

I drafted a fix in #1448 but it still needs quite some work (mostly because it only works with Pydantic v2 at the moment, and relies on using a `Config` instance as input to `generate`, so it won't fix @andreyLetenkovWefox 's example in https://github.com/koxudaxi/datamodel-code-generator/issues/1435#issuecomment-1646081300).

@koxudaxi would you like to see further work in the direction of #1448? or do you have a simpler idea for the fix? It seems that the handling of `base_class` could be simplified in the parser stack, but this looks quite intrusive to me.
> TL;DR: The issue comes from `base_model` incorrectly overriding `RootModel`.

I was also running into this issue, but not in the same way as @andreyLetenkovWefox. Apparently, if you set the `base_class` parameter and use Pydantic V2, it will incorrectly set the Base Model for Root Models to your custom Base Class.

See the following example:

```openapi
openapi: 3.0.0
components:
  schemas:
    CustomId:
      description: My custom ID
      type: string
      format: uuid
```

```bash
datamodel-codegen --input custom_id.yaml --base-class "utils.dto_class.BaseDtoModel" --output-model-type pydantic_v2.BaseModel
```

This will output:

``` python
from __future__ import annotations

from uuid import UUID

from pydantic import Field
from utils.dto_class import BaseDtoModel


class CustomId(BaseDtoModel):
    root: UUID = Field(..., description='My custom ID')
```

This, for Pydantic V2, is wrong as it will be interpreted as a normal model with a literal `root` field (and not a Root Model). In Pydantic V1 this was not an issue because it was being populated through the `__root__` field which uses `BaseModel` as the upper class (just like a normal model).

But here arises another question: **Should `base_class` actually override `RootModel` as well?** In Pydantic V1 this made sense, but for Pydantic V2 probably RootModel should never be overwritten.  
IMO, `RootModel` should not be overwritten even with a custom `base_class` as it implements several methods which are not available in normal models. However, a developer may want to define their custom `RootModel` as well, so not sure what's the correct procedure for this...

Resolving patch:
diff --git a/datamodel_code_generator/model/pydantic_v2/root_model.py b/datamodel_code_generator/model/pydantic_v2/root_model.py
--- a/datamodel_code_generator/model/pydantic_v2/root_model.py
+++ b/datamodel_code_generator/model/pydantic_v2/root_model.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import ClassVar
+from typing import Any, ClassVar
 
 from datamodel_code_generator.model.pydantic_v2.base_model import BaseModel
 
@@ -8,3 +8,14 @@
 class RootModel(BaseModel):
     TEMPLATE_FILE_PATH: ClassVar[str] = 'pydantic_v2/RootModel.jinja2'
     BASE_CLASS: ClassVar[str] = 'pydantic.RootModel'
+
+    def __init__(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        # Remove custom_base_class for Pydantic V2 models; behaviour is different from Pydantic V1 as it will not
+        # be treated as a root model. custom_base_class cannot both implement BaseModel and RootModel!
+        if 'custom_base_class' in kwargs:
+            kwargs.pop('custom_base_class')
+
+        super().__init__(**kwargs)


Test patch:
diff --git a/tests/model/pydantic_v2/__init__.py b/tests/model/pydantic_v2/__init__.py
new file mode 100644
diff --git a/tests/model/pydantic_v2/test_root_model.py b/tests/model/pydantic_v2/test_root_model.py
new file mode 100644
--- /dev/null
+++ b/tests/model/pydantic_v2/test_root_model.py
@@ -0,0 +1,67 @@
+from datamodel_code_generator.model import DataModelFieldBase
+from datamodel_code_generator.model.pydantic_v2.root_model import RootModel
+from datamodel_code_generator.reference import Reference
+from datamodel_code_generator.types import DataType
+
+
+def test_root_model():
+    root_model = RootModel(
+        fields=[
+            DataModelFieldBase(
+                name='a',
+                data_type=DataType(type='str'),
+                default='abc',
+                required=False,
+            )
+        ],
+        reference=Reference(name='TestRootModel', path='test_root_model'),
+    )
+
+    assert root_model.name == 'TestRootModel'
+    assert root_model.fields == [
+        DataModelFieldBase(
+            name='a',
+            data_type=DataType(type='str'),
+            default='abc',
+            required=False,
+            parent=root_model,
+        )
+    ]
+
+    assert root_model.base_class == 'RootModel'
+    assert root_model.custom_base_class is None
+    assert root_model.render() == (
+        'class TestRootModel(RootModel):\n' "    root: Optional[str] = 'abc'"
+    )
+
+
+def test_root_model_custom_base_class():
+    root_model = RootModel(
+        custom_base_class='test.Test',
+        fields=[
+            DataModelFieldBase(
+                name='a',
+                data_type=DataType(type='str'),
+                default='abc',
+                required=False,
+            )
+        ],
+        reference=Reference(name='TestRootModel', path='test_root_model'),
+    )
+
+    assert root_model.name == 'TestRootModel'
+    assert root_model.fields == [
+        DataModelFieldBase(
+            name='a',
+            data_type=DataType(type='str'),
+            default='abc',
+            required=False,
+            parent=root_model,
+        )
+    ]
+
+    assert root_model.base_class == 'RootModel'
+    assert root_model.custom_base_class is None
+    assert root_model.render() == (
+        'class TestRootModel(RootModel):\n' "    root: Optional[str] = 'abc'"
+    )