diff --git a/changes/1581-therefromhere.md b/changes/1581-therefromhere.md new file mode 100644 index 0000000000..389e0af015 --- /dev/null +++ b/changes/1581-therefromhere.md @@ -0,0 +1 @@ +Make `__modify_schema__` on Enums apply to the enum schema rather than fields that use the enum. \ No newline at end of file diff --git a/pydantic/schema.py b/pydantic/schema.py index 3ca516a906..b12c8d0c67 100644 --- a/pydantic/schema.py +++ b/pydantic/schema.py @@ -545,6 +545,10 @@ def enum_process_schema(enum: Type[Enum]) -> Dict[str, Any]: add_field_type_to_schema(enum, schema) + modify_schema = getattr(enum, '__modify_schema__', None) + if modify_schema: + modify_schema(schema) + return schema @@ -698,9 +702,9 @@ def field_singleton_schema( # noqa: C901 (ignore complexity) else: add_field_type_to_schema(field_type, f_schema) - modify_schema = getattr(field_type, '__modify_schema__', None) - if modify_schema: - modify_schema(f_schema) + modify_schema = getattr(field_type, '__modify_schema__', None) + if modify_schema: + modify_schema(f_schema) if f_schema: return f_schema, definitions, nested_models diff --git a/tests/test_schema.py b/tests/test_schema.py index ce447d5ec7..d3326420f6 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -222,6 +222,34 @@ class Model(BaseModel): } +def test_enum_modify_schema(): + class SpamEnum(str, Enum): + foo = 'f' + bar = 'b' + + @classmethod + def __modify_schema__(cls, field_schema): + field_schema['tsEnumNames'] = [e.name for e in cls] + + class Model(BaseModel): + spam: SpamEnum = Field(None) + + assert Model.schema() == { + 'definitions': { + 'SpamEnum': { + 'description': 'An enumeration.', + 'enum': ['f', 'b'], + 'title': 'SpamEnum', + 'tsEnumNames': ['foo', 'bar'], + 'type': 'string', + } + }, + 'properties': {'spam': {'$ref': '#/definitions/SpamEnum'}}, + 'title': 'Model', + 'type': 'object', + } + + def test_json_schema(): class Model(BaseModel): a = b'foobar'