enhance: implement json_schema_extra support for .had() (#112) (#298)

### Summary
This PR fixes the missing support for `json_schema_extra` in the
`.had()` function. Both callable and dict forms are now properly applied
when creating runtime schemas.

### Changes
- `.had()` now correctly applies `json_schema_extra`.
- Unit tests updated to verify this behavior.

### Issue
Closes #112

---------

Co-authored-by: Stanislav Zmiev <[email protected]>

Author: csemanish12

cadwyn/structure/schemas.py

@@ -58,6 +58,7 @@
     "allow_mutation",
     "pattern",
     "discriminator",
+    "json_schema_extra",
 ]
 
 
@@ -100,6 +101,7 @@ class FieldChanges:
     allow_mutation: bool
     pattern: str
     discriminator: str
+    json_schema_extra: Union[dict[str, Any], Callable[[dict[str, Any]], None], None]
 
 
 @dataclass(**DATACLASS_SLOTS)
@@ -178,6 +180,7 @@ def had(
         allow_mutation: bool = Sentinel,
         pattern: str = Sentinel,
         discriminator: str = Sentinel,
+        json_schema_extra: Union[dict[str, Any], Callable[[dict[str, Any]], None], None] = Sentinel,
     ) -> FieldHadInstruction:
         return FieldHadInstruction(
             is_hidden_from_changelog=False,
@@ -222,6 +225,7 @@ def had(
                 allow_mutation=allow_mutation,
                 pattern=pattern,
                 discriminator=discriminator,
+                json_schema_extra=json_schema_extra,
             ),
         )
 

tests/test_schema_generation/test_schema_field.py

@@ -10,6 +10,7 @@
     CadwynStructureError,
     InvalidGenerationInstructionError,
 )
+from cadwyn.schema_generation import SchemaGenerator
 from cadwyn.structure import schema
 from tests.conftest import (
     CreateRuntimeSchemas,
@@ -181,7 +182,7 @@ def assert_field_had_changes_apply(
     attr: str,
     attr_value: Any,
     create_runtime_schemas: CreateRuntimeSchemas,
-):
+) -> dict[str, SchemaGenerator]:
     schemas = create_runtime_schemas(version_change(schema(model).field("foo").had(**{attr: attr_value})))
 
     field_info = schemas["2000-01-01"][model].model_fields["foo"]
@@ -191,6 +192,8 @@ def assert_field_had_changes_apply(
     else:
         assert getattr(field_info, attr) == attr_value
 
+    return schemas
+
 
 @pytest.mark.parametrize(
     ("attr", "attr_value"),
@@ -295,6 +298,46 @@ class SchemaWithOneFloatField(BaseModel):
     )
 
 
+def test_schema_field_had__json_schema_extra_as_dict(create_runtime_schemas: CreateRuntimeSchemas):
+    class SchemaWithFooHadDictJsonSchemaExtra(BaseModel):
+        foo: str
+
+    schemas = assert_field_had_changes_apply(
+        SchemaWithFooHadDictJsonSchemaExtra,
+        "json_schema_extra",
+        {"example": "bar"},
+        create_runtime_schemas,
+    )
+
+    # Validate the JSON Schema produced by Pydantic contains the extra
+    model_cls = schemas["2000-01-01"][SchemaWithFooHadDictJsonSchemaExtra]  # pyright: ignore
+    json_schema = model_cls.model_json_schema()
+    assert json_schema["properties"]["foo"]["example"] == "bar"
+
+
+def test__schema_field_had__json_schema_extra_as_callable(create_runtime_schemas: CreateRuntimeSchemas):
+    def modify_schema(schema_dict: dict[str, Any]):
+        schema_dict["example"] = "bar"
+        schema_dict["custom"] = 42
+
+    class SchemaWithFooHadCallableJsonSchemaExtra(BaseModel):
+        foo: str
+
+    schemas = assert_field_had_changes_apply(
+        SchemaWithFooHadCallableJsonSchemaExtra,
+        "json_schema_extra",
+        modify_schema,
+        create_runtime_schemas,
+    )
+
+    # Validate the JSON Schema produced by Pydantic contains changes from the callable
+    model_cls = schemas["2000-01-01"][SchemaWithFooHadCallableJsonSchemaExtra]  # pyright: ignore
+    json_schema = model_cls.model_json_schema()
+    props = json_schema["properties"]["foo"]
+    assert props["example"] == "bar"
+    assert props["custom"] == 42
+
+
 def test__schema_field_didnt_have__removing_default(create_runtime_schemas: CreateRuntimeSchemas):
     class SchemaWithDefaults(BaseModel):
         foo: str = "hewwo"