Fix typehint for JSON schema extra callable (#6659)

pydantic/config.py

@@ -1,10 +1,10 @@
 """Configuration for Pydantic models."""
 from __future__ import annotations as _annotations
 
-from typing import TYPE_CHECKING, Any, Callable, overload
+from typing import TYPE_CHECKING, Any, Callable, Dict, Type, Union
 from warnings import warn
 
-from typing_extensions import Literal, Protocol, TypedDict
+from typing_extensions import Literal, TypeAlias, TypedDict
 
 from ._migration import getattr_migration
 from .deprecated.config import BaseConfig
@@ -18,14 +18,10 @@
 __all__ = 'BaseConfig', 'ConfigDict', 'Extra'
 
 
-class JsonSchemaExtraCallable(Protocol):
-    @overload
-    def __call__(self, schema: dict[str, Any]) -> None:
-        pass
-
-    @overload
-    def __call__(self, schema: dict[str, Any], model_class: type[Any]) -> None:
-        pass
+JsonSchemaExtraCallable: TypeAlias = Union[
+    Callable[[Dict[str, Any]], None],
+    Callable[[Dict[str, Any], Type[Any]], None],
+]
 
 
 class _Extra:

pydantic/json_schema.py

@@ -1300,9 +1300,9 @@ def _update_class_schema(
             schema_to_update.update(json_schema_extra)
         elif callable(json_schema_extra):
             if len(inspect.signature(json_schema_extra).parameters) > 1:
-                json_schema_extra(schema_to_update, cls)
+                json_schema_extra(schema_to_update, cls)  # type: ignore
             else:
-                json_schema_extra(schema_to_update)
+                json_schema_extra(schema_to_update)  # type: ignore
         elif json_schema_extra is not None:
             raise ValueError(
                 f"model_config['json_schema_extra']={json_schema_extra} should be a dict, callable, or None"