Standadize __get_pydantic_core_schema__ signature

pydantic · Sep 12, 2023 · 6746688 · 6746688
1 parent ed009ba
commit 6746688
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git a/docs/usage/types/custom.md b/docs/usage/types/custom.md
@@ -418,6 +418,7 @@ from typing_extensions import Annotated
 
 from pydantic import (
     BaseModel,
+    GetCoreSchemaHandler,
     GetJsonSchemaHandler,
     ValidationError,
 )
@@ -441,7 +442,7 @@ class _ThirdPartyTypePydanticAnnotation:
     def __get_pydantic_core_schema__(
         cls,
         _source_type: Any,
-        _handler: Callable[[Any], core_schema.CoreSchema],
+        _handler: GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
         """
         We return a pydantic_core.CoreSchema that behaves in the following ways:
@@ -748,7 +749,7 @@ from typing import Any, Callable, Sequence, TypeVar
 from pydantic_core import ValidationError, core_schema
 from typing_extensions import get_args
 
-from pydantic import BaseModel
+from pydantic import BaseModel, GetCoreSchemaHandler
 
 T = TypeVar('T')
 
@@ -765,7 +766,7 @@ class MySequence(Sequence[T]):
 
     @classmethod
     def __get_pydantic_core_schema__(
-        cls, source: Any, handler: Callable[[Any], core_schema.CoreSchema]
+        cls, source: Any, handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
         instance_schema = core_schema.is_instance_schema(cls)
 

diff --git a/pydantic/networks.py b/pydantic/networks.py
@@ -9,7 +9,7 @@
 from pydantic_core import MultiHostUrl, PydanticCustomError, Url, core_schema
 from typing_extensions import Annotated, TypeAlias
 
-from ._internal import _fields, _repr, _schema_generation_shared
+from ._internal import _annotated_handlers, _fields, _repr, _schema_generation_shared
 from ._migration import getattr_migration
 from .json_schema import JsonSchemaValue
 
@@ -194,6 +194,7 @@ class Model(BaseModel):
         def __get_pydantic_core_schema__(
             cls,
             source: type[Any],
+            handler: _annotated_handlers.GetCoreSchemaHandler,
         ) -> core_schema.CoreSchema:
             import_email_validator()
             return core_schema.general_after_validator_function(cls._validate, core_schema.str_schema())
@@ -268,6 +269,7 @@ def __get_pydantic_json_schema__(
     def __get_pydantic_core_schema__(
         cls,
         source: type[Any],
+        handler: _annotated_handlers.GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
         import_email_validator()
         return core_schema.general_after_validator_function(
@@ -321,6 +323,7 @@ def __get_pydantic_json_schema__(
     def __get_pydantic_core_schema__(
         cls,
         source: type[Any],
+        handler: _annotated_handlers.GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
         return core_schema.general_plain_validator_function(
             cls._validate, serialization=core_schema.to_string_ser_schema()
@@ -360,6 +363,7 @@ def __get_pydantic_json_schema__(
     def __get_pydantic_core_schema__(
         cls,
         source: type[Any],
+        handler: _annotated_handlers.GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
         return core_schema.general_plain_validator_function(
             cls._validate, serialization=core_schema.to_string_ser_schema()
@@ -401,6 +405,7 @@ def __get_pydantic_json_schema__(
     def __get_pydantic_core_schema__(
         cls,
         source: type[Any],
+        handler: _annotated_handlers.GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
         return core_schema.general_plain_validator_function(
             cls._validate, serialization=core_schema.to_string_ser_schema()

diff --git a/pydantic/types.py b/pydantic/types.py
@@ -1534,7 +1534,7 @@ def __get_pydantic_json_schema__(
         return field_schema
 
     def __get_pydantic_core_schema__(
-        self, source: type[Any], handler: Callable[[Any], core_schema.CoreSchema]
+        self, source: type[Any], handler: _annotated_handlers.GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
         return core_schema.general_after_validator_function(
             function=self.decode,
@@ -1623,7 +1623,7 @@ class Model(BaseModel):
     """
 
     def __get_pydantic_core_schema__(
-        self, source: type[Any], handler: Callable[[Any], core_schema.CoreSchema]
+        self, source: type[Any], handler: _annotated_handlers.GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
         return core_schema.general_after_validator_function(
             function=self.decode_str,

diff --git a/tests/test_edge_cases.py b/tests/test_edge_cases.py
@@ -29,6 +29,7 @@
 from pydantic import (
     BaseModel,
     ConfigDict,
+    GetCoreSchemaHandler,
     PydanticDeprecatedSince20,
     PydanticInvalidForJsonSchema,
     PydanticSchemaGenerationError,
@@ -1890,10 +1891,7 @@ def __init__(self, t1: T1, t2: T2):
             self.t2 = t2
 
         @classmethod
-        def __get_pydantic_core_schema__(
-            cls,
-            source: Any,
-        ):
+        def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler):
             schema = core_schema.is_instance_schema(cls)
 
             args = get_args(source)