fix: respect --extra-fields option in pydantic v2 models (#2423)

* test --extra-fields with pydantic_v2

* fix how extra fields are handled in pydantic_v2

* refactor to bypass pyright error

* improve test coverage

---------

Co-authored-by: Cosmo Grant <cosmo.grant@postcodelottery.co.uk>

Author: cosmo-grant

src/datamodel_code_generator/model/pydantic_v2/base_model.py

@@ -232,9 +232,20 @@ def __init__(  # noqa: PLR0913
             self.extra_template_data["config"] = ConfigDict.parse_obj(config_parameters)  # pyright: ignore[reportArgumentType]
             self._additional_imports.append(IMPORT_CONFIG_DICT)
 
-    def _get_config_extra(self) -> Literal["'allow'", "'forbid'"] | None:
+    def _get_config_extra(self) -> Literal["'allow'", "'forbid'", "'ignore'"] | None:
         additional_properties = self.extra_template_data.get("additionalProperties")
         allow_extra_fields = self.extra_template_data.get("allow_extra_fields")
-        if additional_properties is not None or allow_extra_fields:
-            return "'allow'" if additional_properties or allow_extra_fields else "'forbid'"
-        return None
+        extra_fields = self.extra_template_data.get("extra_fields")
+
+        config_extra = None
+        if allow_extra_fields or extra_fields == "allow":
+            config_extra = "'allow'"
+        elif extra_fields == "forbid":
+            config_extra = "'forbid'"
+        elif extra_fields == "ignore":
+            config_extra = "'ignore'"
+        elif additional_properties is True:
+            config_extra = "'allow'"
+        elif additional_properties is False:
+            config_extra = "'forbid'"
+        return config_extra

tests/data/expected/main/jsonschema/extra_fields_allow.py

@@ -0,0 +1,39 @@
+# generated by datamodel-codegen:
+#   filename:  extra_fields.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, Extra
+
+
+class Foo(BaseModel):
+    class Config:
+        extra = Extra.allow
+
+    x: Optional[int] = None
+
+
+class Bar(BaseModel):
+    class Config:
+        extra = Extra.allow
+
+    y: Optional[int] = None
+
+
+class Baz(BaseModel):
+    class Config:
+        extra = Extra.allow
+
+    z: Optional[int] = None
+
+
+class Test(BaseModel):
+    class Config:
+        extra = Extra.allow
+
+    foo: Foo
+    bar: Optional[Bar] = None
+    baz: Optional[Baz] = None

tests/data/expected/main/jsonschema/extra_fields_forbid.py

@@ -0,0 +1,39 @@
+# generated by datamodel-codegen:
+#   filename:  extra_fields.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, Extra
+
+
+class Foo(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    x: Optional[int] = None
+
+
+class Bar(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    y: Optional[int] = None
+
+
+class Baz(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    z: Optional[int] = None
+
+
+class Test(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    foo: Foo
+    bar: Optional[Bar] = None
+    baz: Optional[Baz] = None

tests/data/expected/main/jsonschema/extra_fields_ignore.py

@@ -1,5 +1,5 @@
 # generated by datamodel-codegen:
-#   filename:  extra_fields_ignore.json
+#   filename:  extra_fields.json
 #   timestamp: 2019-07-26T00:00:00+00:00
 
 from __future__ import annotations

tests/data/expected/main/jsonschema/extra_fields_v2_allow.py

@@ -0,0 +1,39 @@
+# generated by datamodel-codegen:
+#   filename:  extra_fields.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict
+
+
+class Foo(BaseModel):
+    model_config = ConfigDict(
+        extra='allow',
+    )
+    x: Optional[int] = None
+
+
+class Bar(BaseModel):
+    model_config = ConfigDict(
+        extra='allow',
+    )
+    y: Optional[int] = None
+
+
+class Baz(BaseModel):
+    model_config = ConfigDict(
+        extra='allow',
+    )
+    z: Optional[int] = None
+
+
+class Test(BaseModel):
+    model_config = ConfigDict(
+        extra='allow',
+    )
+    foo: Foo
+    bar: Optional[Bar] = None
+    baz: Optional[Baz] = None

tests/data/expected/main/jsonschema/extra_fields_v2_forbid.py

@@ -0,0 +1,39 @@
+# generated by datamodel-codegen:
+#   filename:  extra_fields.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict
+
+
+class Foo(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    x: Optional[int] = None
+
+
+class Bar(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    y: Optional[int] = None
+
+
+class Baz(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    z: Optional[int] = None
+
+
+class Test(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    foo: Foo
+    bar: Optional[Bar] = None
+    baz: Optional[Baz] = None

tests/data/expected/main/jsonschema/extra_fields_v2_ignore.py

@@ -0,0 +1,39 @@
+# generated by datamodel-codegen:
+#   filename:  extra_fields.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict
+
+
+class Foo(BaseModel):
+    model_config = ConfigDict(
+        extra='ignore',
+    )
+    x: Optional[int] = None
+
+
+class Bar(BaseModel):
+    model_config = ConfigDict(
+        extra='ignore',
+    )
+    y: Optional[int] = None
+
+
+class Baz(BaseModel):
+    model_config = ConfigDict(
+        extra='ignore',
+    )
+    z: Optional[int] = None
+
+
+class Test(BaseModel):
+    model_config = ConfigDict(
+        extra='ignore',
+    )
+    foo: Foo
+    bar: Optional[Bar] = None
+    baz: Optional[Baz] = None

…data/jsonschema/extra_fields_ignore.json tests/data/jsonschema/extra_fields.jsontests/data/jsonschema/extra_fields_ignore.json renamed to tests/data/jsonschema/extra_fields.json tests/data/jsonschema/extra_fields_ignore.json renamed to tests/data/jsonschema/extra_fields.json

@@ -3249,19 +3249,56 @@ def test_main_json_pointer_percent_encoded_segments() -> None:
         assert "".join(result.split()) == "".join(expected.split())
 
 
+@pytest.mark.parametrize(
+    ("extra_fields", "output_model", "expected_output"),
+    [
+        (
+            "allow",
+            "pydantic.BaseModel",
+            "extra_fields_allow.py",
+        ),
+        (
+            "forbid",
+            "pydantic.BaseModel",
+            "extra_fields_forbid.py",
+        ),
+        (
+            "ignore",
+            "pydantic.BaseModel",
+            "extra_fields_ignore.py",
+        ),
+        (
+            "allow",
+            "pydantic_v2.BaseModel",
+            "extra_fields_v2_allow.py",
+        ),
+        (
+            "forbid",
+            "pydantic_v2.BaseModel",
+            "extra_fields_v2_forbid.py",
+        ),
+        (
+            "ignore",
+            "pydantic_v2.BaseModel",
+            "extra_fields_v2_ignore.py",
+        ),
+    ],
+)
 @freeze_time("2019-07-26")
-def test_main_extra_fields_ignore() -> None:
+def test_main_extra_fields(extra_fields: str, output_model: str, expected_output: str) -> None:
     with TemporaryDirectory() as output_dir:
         output_file: Path = Path(output_dir) / "output.py"
         return_code: Exit = main([
             "--input",
-            str(JSON_SCHEMA_DATA_PATH / "extra_fields_ignore.json"),
+            str(JSON_SCHEMA_DATA_PATH / "extra_fields.json"),
             "--output",
             str(output_file),
             "--input-file-type",
             "jsonschema",
             "--extra-fields",
-            "ignore",
+            extra_fields,
+            "--output-model-type",
+            output_model,
         ])
         assert return_code == Exit.OK
-        assert output_file.read_text() == (EXPECTED_JSON_SCHEMA_PATH / "extra_fields_ignore.py").read_text()
+        assert output_file.read_text() == (EXPECTED_JSON_SCHEMA_PATH / expected_output).read_text()