Repository: pydantic/pydantic

Issue:
Possible regression of __modify_schema__ for enum fields
# Bug

A side-effect of #1432 is that `__modify_schema__` on enum subclasses is still affecting the field schema, rather than the enumeration schema.

I'd like it if `__modify_schema__` applied to the resulting enum model.  Does that seem reasonable?

In my use case I'm subclassing enum to add a `tsEnumNames` property, so that I could generate proper typescript enums from the resulting Json schema (using json-schema-to-typescript - see https://github.com/bcherny/json-schema-to-typescript/pull/15 ).

Using latest master 2eb62a3b2f2d345aa7b4625306331f68a7bc0fd5 (or d560c57dbcfe91e9f67782ede3e871ce19b88373 from #1562)


```py
import enum
import pydantic


class TsCompatibleEnum(enum.Enum):
  """Enum subclass that adds tsEnumNames to enumeration fields
  for compatibility with json-schema-to-typescript

  See https://github.com/bcherny/json-schema-to-typescript/pull/15
  """
  @classmethod
  def __modify_schema__(cls, field_schema):
    field_schema["tsEnumNames"] = list(e.name for e in cls)


class MyEnum(TsCompatibleEnum):
  """Enumeration - note that the names and values are different"""
  FOO = "foo"
  BAR = "bar"


class MyModel(pydantic.BaseModel):
  my_field: MyEnum = MyEnum.FOO


print(MyModel.schema_json(indent=2))
```

Outputs:

```json
{
  "title": "MyModel",
  "type": "object",
  "properties": {
    "my_field": {
      "$ref": "#/definitions/MyEnum",
      "tsEnumNames": [
        "FOO",
        "BAR"
      ]
    }
  },
  "definitions": {
    "MyEnum": {
      "title": "MyEnum",
      "description": "Enumeration - note that the names and values are different",
      "enum": [
        "foo",
        "bar"
      ]
    }
  }
}
```

vs with pydantic 1.5.1:

```json
{
  "title": "MyModel",
  "type": "object",
  "properties": {
    "my_field": {
      "title": "My Field",
      "default": "foo",
      "enum": [
        "foo",
        "bar"
      ],
      "tsEnumNames": [
        "FOO",
        "BAR"
      ]
    }
  }
}
```

I'd like to get this instead:

```json
{
  "title": "MyModel",
  "type": "object",
  "properties": {
    "my_field": {
      "$ref": "#/definitions/MyEnum"
    }
  },
  "definitions": {
    "MyEnum": {
      "title": "MyEnum",
      "description": "Enumeration - note that the names and values are different",
      "enum": [
        "foo",
        "bar"
      ],
      "tsEnumNames": [
        "FOO",
        "BAR"
      ]
    }
  }
}
```


Comments:

Patch:
diff --git a/pydantic/schema.py b/pydantic/schema.py
--- 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


Test patch:
diff --git a/tests/test_schema.py b/tests/test_schema.py
--- 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'