Repository: samuelcolvin/pydantic
Pull request #3434 (2021-11-21 20:26:12)

Add optional `field` argument to `__modify_schema__`
<!-- Thank you for your contribution! -->
<!-- Unless your change is trivial, please create an issue to discuss the change before creating a PR -->
<!-- See https://pydantic-docs.helpmanual.io/contributing/ for help on Contributing -->

## Change Summary

Add optional `field` parameter to `model.__modify_schema__()` to opt in to receiving the `ModelField` object of the field. This makes it possible e.g. for the implementation of `__modify_schema__()` of generic type to know the runtime type of its subfields.

The concept is inspired by the optional parameters in validators, implemented in the `pydantic.class_validators` module. There are some differences:

* Only `field` is supported. `values` does not make sense in this context, since validation deals with concrete *object*, and this with the class itself. Model `config` might be supported, but is not included in this PR.
* `field` is `Optional[ModelField]` as opposed to `ModelField`, unlike validators. This is because the `schema` methods can be used on things that are not model fields, in which case `None` will be passed to as the `field` argument.
* Whereas `class_validators` does plenty of checks on the validator method signature before wrapping it, my implementation simply checks for the presence of `field`, and calls `__modify_schema__()` with or without the `field` parameter. This is because with validators there is separate prepare and execute steps, whereas `__modify_schema__()` is called directly when generating schema. If the signature doesn't match the expected, the error message and stack trace will be clear enough.

## Related issue number

There are multiple requests for this:
* https://github.com/samuelcolvin/pydantic/discussions/2902
* https://github.com/samuelcolvin/pydantic/discussions/3445
* My own [hrefs](https://github.com/jasujm/hrefs) project also uses generic types as model fields and needs to access the field spec to properly generate schema

## Checklist

* [X] Unit tests for the changes exist
* [X] Tests pass on CI and coverage remains at 100% (expected to pass since actually running on CI is gate keeped)
* [X] Documentation reflects the changes where applicable
* [X] `changes/<pull request or issue id>-<github username>.md` file added describing change
  (see [changes/README.md](https://github.com/samuelcolvin/pydantic/blob/master/changes/README.md) for details)
* [X] My PR is ready to review, **please add a comment including the phrase "please review" to assign reviewers**


jasujm: please review
samuelcolvin: please update
jasujm: Thank you for the suggestions. I updated and rebased the PR accordingly. please review

Patch:
diff --git a/changes/3175-christianbundy.md b/changes/3175-christianbundy.md
new file mode 100644
index 000000000..3faf22eaf
--- /dev/null
+++ b/changes/3175-christianbundy.md
@@ -0,0 +1,1 @@
+Add support for Mypy 0.920
diff --git a/changes/3434-jasujm.md b/changes/3434-jasujm.md
new file mode 100644
index 000000000..bd93d666f
--- /dev/null
+++ b/changes/3434-jasujm.md
@@ -0,0 +1,2 @@
+When generating field schema, pass optional `field` argument (of type
+`pydantic.fields.ModelField`) to `__modify_schema__()` if present
diff --git a/docs/examples/schema_with_field.py b/docs/examples/schema_with_field.py
new file mode 100644
index 000000000..9288efd8c
--- /dev/null
+++ b/docs/examples/schema_with_field.py
@@ -0,0 +1,31 @@
+# output-json
+from typing import Optional
+
+from pydantic import BaseModel, Field
+from pydantic.fields import ModelField
+
+
+class RestrictedAlphabetStr(str):
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, value, field: ModelField):
+        alphabet = field.field_info.extra['alphabet']
+        if any(c not in alphabet for c in value):
+            raise ValueError(f'{value!r} is not restricted to {alphabet!r}')
+        return cls(value)
+
+    @classmethod
+    def __modify_schema__(cls, field_schema, field: Optional[ModelField]):
+        if field:
+            alphabet = field.field_info.extra['alphabet']
+            field_schema['examples'] = [c * 3 for c in alphabet]
+
+
+class MyModel(BaseModel):
+    value: RestrictedAlphabetStr = Field(alphabet='ABC')
+
+
+print(MyModel.schema_json(indent=2))
diff --git a/docs/requirements.txt b/docs/requirements.txt
index dee12cff6..2bc1b7dc1 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,12 +1,12 @@
 ansi2html==1.6.0
 flake8==4.0.1
 flake8-quotes==3.3.1
-hypothesis==6.30.1
+hypothesis==6.31.4
 markdown-include==0.6.0
 mdx-truly-sane-lists==1.2
 mkdocs==1.2.3
 mkdocs-exclude==1.0.2
-mkdocs-material==8.0.5
+mkdocs-material==8.1.1
 sqlalchemy
 orjson
 ujson
diff --git a/docs/usage/schema.md b/docs/usage/schema.md
index 2ee1a634d..9928b2782 100644
--- a/docs/usage/schema.md
+++ b/docs/usage/schema.md
@@ -150,6 +150,21 @@ For versions of Python prior to 3.9, `typing_extensions.Annotated` can be used.
 Custom field types can customise the schema generated for them using the `__modify_schema__` class method;
 see [Custom Data Types](types.md#custom-data-types) for more details.
 
+`__modify_schema__` can also take a `field` argument which will have type `Optional[ModelField]`.
+*pydantic* will inspect the signature of `__modify_schema__` to determine whether the `field` argument should be
+included.
+
+```py
+{!.tmp_examples/schema_with_field.py!}
+```
+_(This script is complete, it should run "as is")_
+
+Outputs:
+
+```json
+{!.tmp_examples/schema_with_field.json!}
+```
+
 ## JSON Schema Types
 
 Types, custom field types, and constraints (like `max_length`) are mapped to the corresponding spec formats in the
diff --git a/pydantic/mypy.py b/pydantic/mypy.py
index 119c1a07c..c2120992f 100644
--- a/pydantic/mypy.py
+++ b/pydantic/mypy.py
@@ -14,6 +14,14 @@ except ImportError:  # pragma: no cover
 
         warnings.warn('No TOML parser installed, cannot read configuration from `pyproject.toml`.')
         toml = None  # type: ignore
+
+try:
+    from mypy.types import TypeVarDef
+except ImportError:  # pragma: no cover
+    # Backward-compatible with TypeVarDef from Mypy 0.910.
+    from mypy.types import TypeVarType as TypeVarDef  # type: ignore[misc]
+
+
 from mypy.errorcodes import ErrorCode
 from mypy.nodes import (
     ARG_NAMED,
@@ -59,7 +67,6 @@ from mypy.types import (
     Type,
     TypeOfAny,
     TypeType,
-    TypeVarDef,
     TypeVarType,
     UnionType,
     get_proper_type,
@@ -359,7 +366,13 @@ class PydanticModelTransformer:
         tvd = TypeVarDef(self_tvar_name, tvar_fullname, -1, [], obj_type)
         self_tvar_expr = TypeVarExpr(self_tvar_name, tvar_fullname, [], obj_type)
         ctx.cls.info.names[self_tvar_name] = SymbolTableNode(MDEF, self_tvar_expr)
-        self_type = TypeVarType(tvd)
+
+        # Backward-compatible with TypeVarDef from Mypy 0.910.
+        if isinstance(tvd, TypeVarType):  # pragma: no cover
+            self_type = tvd
+        else:
+            self_type = TypeVarType(tvd)
+
         add_method(
             ctx,
             'construct',
diff --git a/pydantic/schema.py b/pydantic/schema.py
index ff6f8af1b..e979678c2 100644
--- a/pydantic/schema.py
+++ b/pydantic/schema.py
@@ -90,6 +90,19 @@ TypeModelOrEnum = Union[Type['BaseModel'], Type[Enum]]
 TypeModelSet = Set[TypeModelOrEnum]
 
 
+def _apply_modify_schema(
+    modify_schema: Callable[..., None], field: Optional[ModelField], field_schema: Dict[str, Any]
+) -> None:
+    from inspect import signature
+
+    sig = signature(modify_schema)
+    args = set(sig.parameters.keys())
+    if 'field' in args or 'kwargs' in args:
+        modify_schema(field_schema, field=field)
+    else:
+        modify_schema(field_schema)
+
+
 def schema(
     models: Sequence[Union[Type['BaseModel'], Type['Dataclass']]],
     *,
@@ -335,7 +348,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
 
 
@@ -567,7 +580,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
 
 
@@ -579,6 +592,7 @@ def model_process_schema(
     ref_prefix: Optional[str] = None,
     ref_template: str = default_ref_template,
     known_models: TypeModelSet = None,
+    field: Optional[ModelField] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]:
     """
     Used by ``model_schema()``, you probably should be using that function.
@@ -592,7 +606,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=field)
         return s, {}, set()
     model = cast(Type['BaseModel'], model)
     s = {'title': model.__config__.title or model.__name__}
@@ -674,7 +688,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: Optional[ModelField] = None) -> Dict[str, Any]:
     """
     Take a single `enum` and generate its schema.
 
@@ -695,7 +709,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_
 
@@ -871,7 +885,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, schema_overrides)
         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=field)
     elif is_namedtuple(field_type):
         sub_schema, *_ = model_process_schema(
             field_type.__pydantic_model__,
@@ -880,6 +894,7 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
             ref_prefix=ref_prefix,
             ref_template=ref_template,
             known_models=known_models,
+            field=field,
         )
         items_schemas = list(sub_schema['properties'].values())
         f_schema.update(
@@ -895,7 +910,7 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
 
         modify_schema = getattr(field_type, '__modify_schema__', None)
         if modify_schema:
-            modify_schema(f_schema)
+            _apply_modify_schema(modify_schema, field, f_schema)
 
     if f_schema:
         return f_schema, definitions, nested_models
@@ -914,6 +929,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


Test patch:
diff --git a/tests/requirements-linting.txt b/tests/requirements-linting.txt
index 1a8712743..81614a46e 100644
--- a/tests/requirements-linting.txt
+++ b/tests/requirements-linting.txt
@@ -1,11 +1,11 @@
 black==21.12b0
 flake8==4.0.1
 flake8-quotes==3.3.1
-hypothesis==6.30.1
+hypothesis==6.31.4
 isort==5.10.1
 mypy==0.910
 pre-commit==2.16.0
 pycodestyle==2.8.0
 pyflakes==2.4.0
-twine==3.7.0
+twine==3.7.1
 types-toml==0.10.1
diff --git a/tests/requirements-testing.txt b/tests/requirements-testing.txt
index cb76137e7..faa291edf 100644
--- a/tests/requirements-testing.txt
+++ b/tests/requirements-testing.txt
@@ -1,5 +1,5 @@
 coverage==6.2
-hypothesis==6.30.1
+hypothesis==6.31.4
 # pin importlib-metadata as upper versions need typing-extensions to work if on python < 3.8
 importlib-metadata==3.1.0;python_version<"3.8"
 mypy==0.910
diff --git a/tests/test_schema.py b/tests/test_schema.py
index c1e1f01b0..eb1f081e8 100644
--- a/tests/test_schema.py
+++ b/tests/test_schema.py
@@ -32,6 +32,7 @@ from typing_extensions import Annotated, Literal
 from pydantic import BaseModel, Extra, Field, ValidationError, confrozenset, conlist, conset, validator
 from pydantic.color import Color
 from pydantic.dataclasses import dataclass
+from pydantic.fields import ModelField
 from pydantic.generics import GenericModel
 from pydantic.networks import AnyUrl, EmailStr, IPvAnyAddress, IPvAnyInterface, IPvAnyNetwork, NameEmail, stricturl
 from pydantic.schema import (
@@ -2628,6 +2629,27 @@ def test_complex_nested_generic():
     }
 
 
+def test_schema_with_field_parameter():
+    class RestrictedAlphabetStr(str):
+        @classmethod
+        def __modify_schema__(cls, field_schema, field: Optional[ModelField]):
+            assert isinstance(field, ModelField)
+            alphabet = field.field_info.extra['alphabet']
+            field_schema['examples'] = [c * 3 for c in alphabet]
+
+    class MyModel(BaseModel):
+        value: RestrictedAlphabetStr = Field(alphabet='ABC')
+
+    assert MyModel.schema() == {
+        'title': 'MyModel',
+        'type': 'object',
+        'properties': {
+            'value': {'title': 'Value', 'alphabet': 'ABC', 'examples': ['AAA', 'BBB', 'CCC'], 'type': 'string'}
+        },
+        'required': ['value'],
+    }
+
+
 def test_discriminated_union():
     class BlackCat(BaseModel):
         pet_type: Literal['cat']