Traverse function-before schemas during schema gathering

Viicos · Viicos · commit fc521388f212 · 2025-04-29T13:33:49.000+02:00
Backport of: #11801
diff --git a/pydantic/_internal/_schema_gather.py b/pydantic/_internal/_schema_gather.py
@@ -161,6 +161,11 @@ def traverse_schema(schema: AllSchemas, context: GatherContext) -> None:
             traverse_schema(schema['return_schema'], context)
     elif schema_type == 'computed-field':
         traverse_schema(schema['return_schema'], context)
+    elif schema_type == 'function-before':
+        if 'schema' in schema:
+            traverse_schema(schema['schema'], context)
+        if 'json_schema_input_schema' in schema:
+            traverse_schema(schema['json_schema_input_schema'], context)
     elif schema_type == 'function-plain':
         # TODO duplicate schema types for serializers and validators, needs to be deduplicated.
         if 'return_schema' in schema:
diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py
@@ -6589,6 +6589,19 @@ class Model(BaseModel):
     }
 
 
+def test_json_schema_input_type_inlined() -> None:
+    class Sub(BaseModel):
+        pass
+
+    class Model(BaseModel):
+        sub: Annotated[object, BeforeValidator(lambda v: v, json_schema_input_type=Sub)]
+
+    json_schema = Model.model_json_schema()
+
+    assert 'Sub' in json_schema['$defs']
+    assert json_schema['properties']['sub'] == {'$ref': '#/$defs/Sub'}
+
+
 @pytest.mark.parametrize(
     'validator',
     [
