Repository: pydantic/pydantic

Issue:
Support passing model class to schema_extra?
# Feature Request

As a follow-up to #1054, it would be useful if the model class was passed to the schema_extra callable.

I think @dmontagu was thinking along the same lines in this comment: https://github.com/samuelcolvin/pydantic/issues/937#issuecomment-546459285 was there a decision not to implement this or would a PR be accepted?

Please complete:
* OS: Ubuntu 18.04
* Python version `import sys; print(sys.version)`: 3.7.5
* Pydantic version `import pydantic; print(pydantic.VERSION)`: 1.3

My use case is that I'm adding sub-class configuration in my own nested `Meta` classes, and I'd like to extract information from them in the base class's `schema_extra`

```py
from typing import Any, Dict

import pydantic


class MyBase(pydantic.BaseModel):
    class Config:
        @staticmethod
        def schema_extra(schema: Dict[str, Any], model_class):
            # ...do something with model_class
            schema["something"] = model_class.MyMeta.foo


class MySubclass(MyBase):
    class MyMeta:
        foo = "bar"
```

edit: I originally suggested allowing schema_extra to be a classmethod, but that doesn't work because the cls param would be Config, not the outer model class.


Discussion:
agreed this would be useful, it was an oversight not to add a model argument when we first implemented this.

Options:
* wait until v2 to add this
* add a `model` argument and break `v1.3` - since this was only released yesterday I think we could call it a patch?
* inspect `schema_extra` and pass `model` only if the argument is expected. This would be entirely backwards compatible but involve more code

Thoughts?
I'd say option 2 or 3?
Let's do option 2. PR welcome.

Resolving patch:
diff --git a/docs/examples/schema_extra_callable.py b/docs/examples/schema_extra_callable.py
--- a/docs/examples/schema_extra_callable.py
+++ b/docs/examples/schema_extra_callable.py
@@ -1,5 +1,5 @@
 # output-json
-from typing import Dict, Any
+from typing import Dict, Any, Type
 from pydantic import BaseModel
 
 class Person(BaseModel):
@@ -8,7 +8,7 @@ class Person(BaseModel):
 
     class Config:
         @staticmethod
-        def schema_extra(schema: Dict[str, Any]) -> None:
+        def schema_extra(schema: Dict[str, Any], model: Type['Person']) -> None:
             for prop in schema.get('properties', {}).values():
                 prop.pop('title', None)
 
diff --git a/pydantic/schema.py b/pydantic/schema.py
--- a/pydantic/schema.py
+++ b/pydantic/schema.py
@@ -5,6 +5,7 @@
 from enum import Enum
 from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network
 from pathlib import Path
+from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -450,7 +451,7 @@ def model_process_schema(
     sub-models of the returned schema will be referenced, but their definitions will not be included in the schema. All
     the definitions are returned as the second value.
     """
-    from inspect import getdoc
+    from inspect import getdoc, signature
 
     ref_prefix = ref_prefix or default_prefix
     known_models = known_models or set()
@@ -465,7 +466,12 @@ def model_process_schema(
     s.update(m_schema)
     schema_extra = model.__config__.schema_extra
     if callable(schema_extra):
-        schema_extra(s)
+        if not isinstance(schema_extra, FunctionType):
+            raise TypeError(f'{model.__name__}.Config.schema_extra callable is expected to be a staticmethod')
+        if len(signature(schema_extra).parameters) == 1:
+            schema_extra(s)
+        else:
+            schema_extra(s, model)
     else:
         s.update(schema_extra)
     return s, m_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
@@ -1490,6 +1490,20 @@ class Config:
 
 
 def test_model_with_schema_extra_callable():
+    class Model(BaseModel):
+        name: str = None
+
+        class Config:
+            @staticmethod
+            def schema_extra(schema, model_class):
+                schema.pop('properties')
+                schema['type'] = 'override'
+                assert model_class is Model
+
+    assert Model.schema() == {'title': 'Model', 'type': 'override'}
+
+
+def test_model_with_schema_extra_callable_no_model_class():
     class Model(BaseModel):
         name: str = None
 
@@ -1502,6 +1516,20 @@ def schema_extra(schema):
     assert Model.schema() == {'title': 'Model', 'type': 'override'}
 
 
+def test_model_with_schema_extra_callable_classmethod_asserts():
+    class Model(BaseModel):
+        name: str = None
+
+        class Config:
+            @classmethod
+            def schema_extra(cls, schema, model_class):
+                schema.pop('properties')
+                schema['type'] = 'override'
+
+    with pytest.raises(TypeError, match='Model.Config.schema_extra callable is expected to be a staticmethod'):
+        Model.schema()
+
+
 def test_model_with_extra_forbidden():
     class Model(BaseModel):
         a: str