1"""Types and utility functions used by various other internal tools."""
2
3from __future__ import annotations
4
5from typing import TYPE_CHECKING, Any, Callable, Literal
6
7from pydantic_core import core_schema
8
9from ..annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler
10
11if TYPE_CHECKING:
12 from ..json_schema import GenerateJsonSchema, JsonSchemaValue
13 from ._core_utils import CoreSchemaOrField
14 from ._generate_schema import GenerateSchema
15 from ._namespace_utils import NamespacesTuple
16
17 GetJsonSchemaFunction = Callable[[CoreSchemaOrField, GetJsonSchemaHandler], JsonSchemaValue]
18 HandlerOverride = Callable[[CoreSchemaOrField], JsonSchemaValue]
19
20
21class GenerateJsonSchemaHandler(GetJsonSchemaHandler):
22 """JsonSchemaHandler implementation that doesn't do ref unwrapping by default.
23
24 This is used for any Annotated metadata so that we don't end up with conflicting
25 modifications to the definition schema.
26
27 Used internally by Pydantic, please do not rely on this implementation.
28 See `GetJsonSchemaHandler` for the handler API.
29 """
30
31 def __init__(self, generate_json_schema: GenerateJsonSchema, handler_override: HandlerOverride | None) -> None:
32 self.generate_json_schema = generate_json_schema
33 self.handler = handler_override or generate_json_schema.generate_inner
34 self.mode = generate_json_schema.mode
35
36 def __call__(self, core_schema: CoreSchemaOrField, /) -> JsonSchemaValue:
37 return self.handler(core_schema)
38
39 def resolve_ref_schema(self, maybe_ref_json_schema: JsonSchemaValue) -> JsonSchemaValue:
40 """Resolves `$ref` in the json schema.
41
42 This returns the input json schema if there is no `$ref` in json schema.
43
44 Args:
45 maybe_ref_json_schema: The input json schema that may contains `$ref`.
46
47 Returns:
48 Resolved json schema.
49
50 Raises:
51 LookupError: If it can't find the definition for `$ref`.
52 """
53 if '$ref' not in maybe_ref_json_schema:
54 return maybe_ref_json_schema
55 ref = maybe_ref_json_schema['$ref']
56 json_schema = self.generate_json_schema.get_schema_from_definitions(ref)
57 if json_schema is None:
58 raise LookupError(
59 f'Could not find a ref for {ref}.'
60 ' Maybe you tried to call resolve_ref_schema from within a recursive model?'
61 )
62 return json_schema
63
64
65class CallbackGetCoreSchemaHandler(GetCoreSchemaHandler):
66 """Wrapper to use an arbitrary function as a `GetCoreSchemaHandler`.
67
68 Used internally by Pydantic, please do not rely on this implementation.
69 See `GetCoreSchemaHandler` for the handler API.
70 """
71
72 def __init__(
73 self,
74 handler: Callable[[Any], core_schema.CoreSchema],
75 generate_schema: GenerateSchema,
76 ref_mode: Literal['to-def', 'unpack'] = 'to-def',
77 ) -> None:
78 self._handler = handler
79 self._generate_schema = generate_schema
80 self._ref_mode = ref_mode
81
82 def __call__(self, source_type: Any, /) -> core_schema.CoreSchema:
83 schema = self._handler(source_type)
84 if self._ref_mode == 'to-def':
85 ref = schema.get('ref')
86 if ref is not None:
87 return self._generate_schema.defs.create_definition_reference_schema(schema)
88 return schema
89 else: # ref_mode = 'unpack'
90 return self.resolve_ref_schema(schema)
91
92 def _get_types_namespace(self) -> NamespacesTuple:
93 return self._generate_schema._types_namespace
94
95 def generate_schema(self, source_type: Any, /) -> core_schema.CoreSchema:
96 return self._generate_schema.generate_schema(source_type)
97
98 @property
99 def field_name(self) -> str | None:
100 return self._generate_schema.field_name_stack.get()
101
102 def resolve_ref_schema(self, maybe_ref_schema: core_schema.CoreSchema) -> core_schema.CoreSchema:
103 """Resolves reference in the core schema.
104
105 Args:
106 maybe_ref_schema: The input core schema that may contains reference.
107
108 Returns:
109 Resolved core schema.
110
111 Raises:
112 LookupError: If it can't find the definition for reference.
113 """
114 if maybe_ref_schema['type'] == 'definition-ref':
115 ref = maybe_ref_schema['schema_ref']
116 definition = self._generate_schema.defs.get_schema_from_ref(ref)
117 if definition is None:
118 raise LookupError(
119 f'Could not find a ref for {ref}.'
120 ' Maybe you tried to call resolve_ref_schema from within a recursive model?'
121 )
122 return definition
123 elif maybe_ref_schema['type'] == 'definitions':
124 return self.resolve_ref_schema(maybe_ref_schema['schema'])
125 return maybe_ref_schema