Coverage for /pythoncovmergedfiles/medio/medio/src/pydantic/pydantic/_internal/_json_schema_shared.py: 46%

1""" 

2Types and utility functions used by various other internal tools. 

3""" 

4from __future__ import annotations 

5 

6from typing import TYPE_CHECKING, Any, Callable, Dict, Union 

7 

8from pydantic_core import core_schema 

9 

10if TYPE_CHECKING: 

11 from pydantic.json_schema import GenerateJsonSchema 

12 

13 

14JsonSchemaValue = Dict[str, Any] 

15CoreSchemaOrField = Union[core_schema.CoreSchema, core_schema.DataclassField, core_schema.TypedDictField] 

16 

17 

18class GetJsonSchemaHandler: 

19 """ 

20 Handler to call into the next JSON schema generation function 

21 """ 

22 

23 def __call__(self, __core_schema: CoreSchemaOrField) -> JsonSchemaValue: 

24 """Call the inner handler and get the JsonSchemaValue it returns. 

25 This will call the next JSON schema modifying function up until it calls 

26 into `pydantic.json_schema.GenerateJsonSchema`, which will raise a 

27 `pydantic.errors.PydanticInvalidForJsonSchema` error if it cannot generate 

28 a JSON schema. 

29 

30 Args: 

31 __core_schema (CoreSchemaOrField): A `pydantic_core.core_schema.CoreSchema`. 

32 

33 Returns: 

34 JsonSchemaValue: the JSON schema generated by the inner JSON schema modify 

35 functions. 

36 """ 

37 raise NotImplementedError 

38 

39 def resolve_ref_schema(self, __maybe_ref_json_schema: JsonSchemaValue) -> JsonSchemaValue: 

40 """Get the real schema for a `{"$ref": ...}` schema. 

41 If the schema given is not a `$ref` schema, it will be returned as is. 

42 This means you don't have to check before calling this function. 

43 

44 Args: 

45 __maybe_ref_json_schema (JsonSchemaValue): A JsonSchemaValue, ref based or not. 

46 

47 Raises: 

48 LookupError: if the ref is not found. 

49 

50 Returns: 

51 JsonSchemaValue: A JsonSchemaValue that has no `$ref`. 

52 """ 

53 raise NotImplementedError 

54 

55 

56GetJsonSchemaFunction = Callable[[CoreSchemaOrField, GetJsonSchemaHandler], JsonSchemaValue] 

57 

58 

59class UnpackedRefJsonSchemaHandler(GetJsonSchemaHandler): 

60 """ 

61 A GetJsonSchemaHandler implementation that automatically unpacks `$ref` 

62 schemas so that the caller doesn't have to worry about that. 

63 

64 This is used for custom types and models that implement `__get_pydantic_core_schema__` 

65 so they they always get a `non-$ref` schema. 

66 """ 

67 

68 original_schema: JsonSchemaValue | None = None 

69 

70 def __init__(self, handler: GetJsonSchemaHandler) -> None: 

71 self.handler = handler 

72 

73 def resolve_ref_schema(self, __maybe_ref_json_schema: JsonSchemaValue) -> JsonSchemaValue: 

74 return self.handler.resolve_ref_schema(__maybe_ref_json_schema) 

75 

76 def __call__(self, __core_schema: CoreSchemaOrField) -> JsonSchemaValue: 

77 self.original_schema = self.handler(__core_schema) 

78 return self.resolve_ref_schema(self.original_schema) 

79 

80 def update_schema(self, schema: JsonSchemaValue) -> JsonSchemaValue: 

81 if self.original_schema is None: 

82 # handler / our __call__ was never called 

83 return schema 

84 original_schema = self.resolve_ref_schema(self.original_schema) 

85 if original_schema is not self.original_schema and schema is not original_schema: 

86 # a new schema was returned 

87 original_schema.clear() 

88 original_schema.update(schema) 

89 # return self.original_schema, which may be a ref schema 

90 return self.original_schema 

91 

92 

93def wrap_json_schema_fn_for_model_or_custom_type_with_ref_unpacking( 

94 fn: GetJsonSchemaFunction, 

95) -> GetJsonSchemaFunction: 

96 def wrapped(schema_or_field: CoreSchemaOrField, handler: GetJsonSchemaHandler) -> JsonSchemaValue: 

97 wrapped_handler = UnpackedRefJsonSchemaHandler(handler) 

98 json_schema = fn(schema_or_field, wrapped_handler) 

99 json_schema = wrapped_handler.update_schema(json_schema) 

100 return json_schema 

101 

102 return wrapped 

103 

104 

105HandlerOverride = Callable[[CoreSchemaOrField], JsonSchemaValue] 

106 

107 

108class GenerateJsonSchemaHandler(GetJsonSchemaHandler): 

109 """ 

110 JsonSchemaHandler implementation that doesn't do ref unwrapping by default. 

111 

112 This is used for any Annotated metadata so that we don't end up with conflicting 

113 modifications to the definition schema. 

114 """ 

115 

116 def __init__(self, generate_json_schema: GenerateJsonSchema, handler_override: HandlerOverride | None) -> None: 

117 self.generate_json_schema = generate_json_schema 

118 self.handler = handler_override or generate_json_schema.generate_inner 

119 

120 def __call__(self, __core_schema: CoreSchemaOrField) -> JsonSchemaValue: 

121 return self.handler(__core_schema) 

122 

123 def resolve_ref_schema(self, __maybe_ref_json_schema: JsonSchemaValue) -> JsonSchemaValue: 

124 if '$ref' not in __maybe_ref_json_schema: 

125 return __maybe_ref_json_schema 

126 ref = __maybe_ref_json_schema['$ref'] 

127 json_schema = self.generate_json_schema.get_schema_from_definitions(ref) 

128 if json_schema is None: 

129 raise LookupError( 

130 f'Could not find a ref for {ref}.' 

131 ' Maybe you tried to call resolve_ref_schema from within a recursive model?' 

132 ) 

133 return json_schema