WIP: Add optional field argument to __modify_schema__

pydantic/schema.py

@@ -87,6 +87,16 @@
 TypeModelSet = Set[TypeModelOrEnum]
 
 
+def _apply_modify_schema(modify_schema: Callable, field: ModelField, schema: Dict[str, Any]):
+    from inspect import signature
+
+    sig = signature(modify_schema)
+    if "field" in list(sig.parameters.keys()):
+        modify_schema(schema, field=field)
+    else:
+        modify_schema(schema)
+
+
 def schema(
     models: Sequence[Union[Type['BaseModel'], Type['Dataclass']]],
     *,
@@ -302,7 +312,7 @@ def get_field_schema_validations(field: ModelField) -> Dict[str, Any]:
         f_schema.update(field.field_info.extra)
     modify_schema = getattr(field.outer_type_, '__modify_schema__', None)
     if modify_schema:
-        modify_schema(f_schema)
+        _apply_modify_schema(modify_schema, field, f_schema)
     return f_schema
 
 
@@ -530,7 +540,7 @@ def field_type_schema(
             field_type = field.outer_type_
         modify_schema = getattr(field_type, '__modify_schema__', None)
         if modify_schema:
-            modify_schema(f_schema)
+            _apply_modify_schema(modify_schema, field, f_schema)
     return f_schema, definitions, nested_models
 
 
@@ -542,6 +552,7 @@ def model_process_schema(
     ref_prefix: Optional[str] = None,
     ref_template: str = default_ref_template,
     known_models: TypeModelSet = None,
+    field: ModelField = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]:
     """
     Used by ``model_schema()``, you probably should be using that function.
@@ -555,7 +566,7 @@ def model_process_schema(
     known_models = known_models or set()
     if lenient_issubclass(model, Enum):
         model = cast(Type[Enum], model)
-        s = enum_process_schema(model)
+        s = enum_process_schema(model, field)
         return s, {}, set()
     model = cast(Type['BaseModel'], model)
     s = {'title': model.__config__.title or model.__name__}
@@ -637,7 +648,7 @@ def model_type_schema(
     return out_schema, definitions, nested_models
 
 
-def enum_process_schema(enum: Type[Enum]) -> Dict[str, Any]:
+def enum_process_schema(enum: Type[Enum], field: ModelField) -> Dict[str, Any]:
     """
     Take a single `enum` and generate its schema.
 
@@ -658,7 +669,7 @@ def enum_process_schema(enum: Type[Enum]) -> Dict[str, Any]:
 
     modify_schema = getattr(enum, '__modify_schema__', None)
     if modify_schema:
-        modify_schema(schema_)
+        _apply_modify_schema(modify_schema, field, schema_)
 
     return schema_
 
@@ -834,7 +845,7 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
         enum_name = model_name_map[field_type]
         f_schema, schema_overrides = get_field_info_schema(field)
         f_schema.update(get_schema_ref(enum_name, ref_prefix, ref_template, schema_overrides))
-        definitions[enum_name] = enum_process_schema(field_type)
+        definitions[enum_name] = enum_process_schema(field_type, field)
     elif is_namedtuple(field_type):
         sub_schema, *_ = model_process_schema(
             field_type.__pydantic_model__,
@@ -843,6 +854,7 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
             ref_prefix=ref_prefix,
             ref_template=ref_template,
             known_models=known_models,
+            field=field,
         )
         f_schema.update({'type': 'array', 'items': list(sub_schema['properties'].values())})
     elif not hasattr(field_type, '__pydantic_model__'):
@@ -869,6 +881,7 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
                 ref_prefix=ref_prefix,
                 ref_template=ref_template,
                 known_models=known_models,
+                field=field,
             )
             definitions.update(sub_definitions)
             definitions[model_name] = sub_schema