🐛 Support defining json_schema_extra on RootModel using Field (#…

…6622)

Co-authored-by: David Montague <35119617+dmontagu@users.noreply.github.com>

pydantic/json_schema.py

@@ -1256,6 +1256,16 @@ def model_schema(self, schema: core_schema.ModelSchema) -> JsonSchemaValue:
         title = config.get('title')
 
         json_schema_extra = config.get('json_schema_extra')
+        if cls.__pydantic_root_model__:
+            root_json_schema_extra = cls.model_fields['root'].json_schema_extra
+            if json_schema_extra and root_json_schema_extra:
+                raise ValueError(
+                    '"model_config[\'json_schema_extra\']" and "Field.json_schema_extra" on "RootModel.root"'
+                    ' field must not be set simultaneously'
+                )
+            if root_json_schema_extra:
+                json_schema_extra = root_json_schema_extra
+
         json_schema = self._update_class_schema(json_schema, title, config.get('extra', None), cls, json_schema_extra)
 
         return json_schema

tests/test_root_model.py

@@ -561,3 +561,35 @@ class MyRootModel(RootModel[str]):
 
     MyRootModel = module.MyRootModel
     assert MyRootModel(root='abc') == pickle.loads(pickle.dumps(MyRootModel(root='abc')))
+
+
+def test_json_schema_extra_on_model():
+    class Model(RootModel):
+        model_config = ConfigDict(json_schema_extra={'schema key': 'schema value'})
+        root: str
+
+    assert Model.model_json_schema() == {
+        'schema key': 'schema value',
+        'title': 'Model',
+        'type': 'string',
+    }
+
+
+def test_json_schema_extra_on_field():
+    class Model(RootModel):
+        root: str = Field(json_schema_extra={'schema key': 'schema value'})
+
+    assert Model.model_json_schema() == {
+        'schema key': 'schema value',
+        'title': 'Model',
+        'type': 'string',
+    }
+
+
+def test_json_schema_extra_on_model_and_on_field():
+    class Model(RootModel):
+        model_config = ConfigDict(json_schema_extra={'schema key on model': 'schema value on model'})
+        root: str = Field(json_schema_extra={'schema key on field': 'schema value on field'})
+
+    with pytest.raises(ValueError, match=r'json_schema_extra.*?must not be set simultaneously'):
+        Model.model_json_schema()