Preserve extras in oneOf/anyOf structures during initialization of JsonSchemaObject (#2623)

Author: koxudaxi

src/datamodel_code_generator/parser/jsonschema.py

@@ -341,14 +341,17 @@ class Config:
             keep_untouched = (cached_property,)
             smart_casts = True
 
-    if not TYPE_CHECKING:
-
-        def __init__(self, **data: Any) -> None:
-            """Initialize JsonSchemaObject with extra fields handling."""
-            super().__init__(**data)
-            self.extras = {k: v for k, v in data.items() if k not in EXCLUDE_FIELD_KEYS}
-            if "const" in data.get(self.__extra_key__, {}):  # pragma: no cover
-                self.extras["const"] = data[self.__extra_key__]["const"]
+    def __init__(self, **data: Any) -> None:
+        """Initialize JsonSchemaObject with extra fields handling."""
+        super().__init__(**data)
+        # Restore extras from alias key (for dict -> parse_obj round-trip)
+        alias_extras = data.get(self.__extra_key__, {})
+        # Collect custom keys from raw data
+        raw_extras = {k: v for k, v in data.items() if k not in EXCLUDE_FIELD_KEYS}
+        # Merge: raw_extras takes precedence (original data is the source of truth)
+        self.extras = {**alias_extras, **raw_extras}
+        if "const" in alias_extras:  # pragma: no cover
+            self.extras["const"] = alias_extras["const"]
 
     @cached_property
     def is_object(self) -> bool:

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

@@ -0,0 +1,45 @@
+# generated by datamodel-codegen:
+#   filename:  extras_in_oneof.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class OneofProp(BaseModel):
+    shared_prop: Optional[str] = Field(
+        None, json_schema_extra={'x-shared': 'shared_value'}
+    )
+    variant_a_prop: Optional[str] = Field(
+        None, json_schema_extra={'x-variant': 'variant_a_value'}
+    )
+
+
+class OneofProp1(BaseModel):
+    shared_prop: Optional[str] = Field(
+        None, json_schema_extra={'x-shared': 'shared_value'}
+    )
+    variant_b_prop: Optional[int] = Field(
+        None, json_schema_extra={'x-variant': 'variant_b_value'}
+    )
+
+
+class AnyofProp(BaseModel):
+    any_a_prop: Optional[str] = Field(None, json_schema_extra={'x-any': 'any_a_value'})
+
+
+class AnyofProp1(BaseModel):
+    any_b_prop: Optional[bool] = Field(None, json_schema_extra={'x-any': 'any_b_value'})
+
+
+class ExtrasInOneOf(BaseModel):
+    simple_prop: Optional[str] = Field(
+        None, json_schema_extra={'x-custom': 'simple_value'}
+    )
+    oneof_prop: Optional[Union[OneofProp, OneofProp1]] = Field(
+        None, json_schema_extra={'x-parent-custom': 'parent_value'}
+    )
+    anyof_prop: Optional[Union[AnyofProp, AnyofProp1]] = None

tests/data/jsonschema/extras_in_oneof.json

@@ -0,0 +1,49 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "id": "extras_in_oneof",
+    "title": "ExtrasInOneOf",
+    "type": "object",
+    "properties": {
+        "simple_prop": {"type": "string", "x-custom": "simple_value"},
+        "oneof_prop": {
+            "x-parent-custom": "parent_value",
+            "properties": {
+                "shared_prop": {"type": "string", "x-shared": "shared_value"}
+            },
+            "oneOf": [
+                {
+                    "title": "VariantA",
+                    "type": "object",
+                    "properties": {
+                        "variant_a_prop": {"type": "string", "x-variant": "variant_a_value"}
+                    }
+                },
+                {
+                    "title": "VariantB",
+                    "type": "object",
+                    "properties": {
+                        "variant_b_prop": {"type": "integer", "x-variant": "variant_b_value"}
+                    }
+                }
+            ]
+        },
+        "anyof_prop": {
+            "anyOf": [
+                {
+                    "title": "AnyVariantA",
+                    "type": "object",
+                    "properties": {
+                        "any_a_prop": {"type": "string", "x-any": "any_a_value"}
+                    }
+                },
+                {
+                    "title": "AnyVariantB",
+                    "type": "object",
+                    "properties": {
+                        "any_b_prop": {"type": "boolean", "x-any": "any_b_value"}
+                    }
+                }
+            ]
+        }
+    }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -3294,3 +3294,20 @@ def test_main_jsonschema_root_model_default_value_non_root(output_file: Path) ->
             "--use-annotated",
         ],
     )
+
+
+@pytest.mark.benchmark
+def test_main_jsonschema_extras_in_oneof(output_file: Path) -> None:
+    """Test that extras are preserved in oneOf/anyOf structures (Issue #2403)."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "extras_in_oneof.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        expected_file="extras_in_oneof.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--field-include-all-keys",
+        ],
+    )