Merge pull request #88 from adrienDog/patch-1

feat: Allow clients to pass json_schema_extra too

src/patito/pydantic.py

@@ -1331,8 +1331,11 @@ def Field(
                     f"unexpected kwarg {kwarg}={kwargs[kwarg]}.  Add modern_kwargs_only=False to ignore"
                 )
     ci_json = ci.model_dump_json()
+    existing_json_schema_extra = kwargs.pop("json_schema_extra", {})
+    merged_json_schema_extra = {**existing_json_schema_extra, "column_info": ci_json}
+
     return fields.Field(
         *args,
-        json_schema_extra={"column_info": ci_json},
+        json_schema_extra=merged_json_schema_extra,
         **kwargs,
     )

tests/test_model.py

@@ -587,3 +587,28 @@ class SingleColumnModel(pt.Model):
     assert len(full_list) == len(df)
     for model_instance in full_list:
         assert isinstance(model_instance, SingleColumnModel)
+
+
+def test_json_schema_extra_is_extended_when_it_exists() -> None:
+    """Ensure that the json_schema_extra property is extended with column_info when it is set from the model field."""
+
+    class Model(pt.Model):
+        a: int
+        b: int = pt.Field(
+            json_schema_extra={"client_column_metadata": {"group1": "x", "group2": "y"}}
+        )
+        c: int = pt.Field(
+            json_schema_extra={"client_column_metadata": {"group1": "xxx"}}
+        )
+
+    schema = Model.model_json_schema()  # no serialization issues
+    props = schema[
+        "properties"
+    ]  # extra fields are stored in modified schema_properties
+    for col in ["b", "c"]:
+        assert "column_info" in props[col]
+        assert "client_column_metadata" in props[col]
+    assert "client_column_metadata" not in props["a"]
+    assert props["b"]["client_column_metadata"]["group1"] == "x"
+    assert props["b"]["client_column_metadata"]["group2"] == "y"
+    assert props["c"]["client_column_metadata"]["group1"] == "xxx"