Skip to content

Commit

Permalink
Apply __modify_schema__ on enum schema rather than fields that use it
Browse files Browse the repository at this point in the history
Resolves #1576
  • Loading branch information
therefromhere committed May 31, 2020
1 parent 2eb62a3 commit 6a66c71
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 3 deletions.
1 change: 1 addition & 0 deletions changes/1581-therefromhere.md
@@ -0,0 +1 @@
Make `__modify_schema__` on Enums apply to the enum schema rather than fields that use the enum.
10 changes: 7 additions & 3 deletions pydantic/schema.py
Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions tests/test_schema.py
Expand Up @@ -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'] = list(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'
Expand Down

0 comments on commit 6a66c71

Please sign in to comment.