Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import inspect
4from collections.abc import Mapping, Sequence
5from typing import TYPE_CHECKING, Any, TypeAlias, TypeGuard
7from pydantic_core import CoreSchema, core_schema
9if TYPE_CHECKING:
10 from rich.console import Console
12AnyFunctionSchema: TypeAlias = (
13 core_schema.AfterValidatorFunctionSchema
14 | core_schema.BeforeValidatorFunctionSchema
15 | core_schema.WrapValidatorFunctionSchema
16 | core_schema.PlainValidatorFunctionSchema
17)
20FunctionSchemaWithInnerSchema: TypeAlias = (
21 core_schema.AfterValidatorFunctionSchema
22 | core_schema.BeforeValidatorFunctionSchema
23 | core_schema.WrapValidatorFunctionSchema
24)
27CoreSchemaField: TypeAlias = (
28 core_schema.ModelField | core_schema.DataclassField | core_schema.TypedDictField | core_schema.ComputedField
29)
31CoreSchemaOrField: TypeAlias = core_schema.CoreSchema | CoreSchemaField
33_CORE_SCHEMA_FIELD_TYPES = {'typed-dict-field', 'dataclass-field', 'model-field', 'computed-field'}
34_FUNCTION_WITH_INNER_SCHEMA_TYPES = {'function-before', 'function-after', 'function-wrap'}
35_LIST_LIKE_SCHEMA_WITH_ITEMS_TYPES = {'list', 'deque', 'set', 'frozenset'}
38def is_core_schema(
39 schema: CoreSchemaOrField,
40) -> TypeGuard[CoreSchema]:
41 return schema['type'] not in _CORE_SCHEMA_FIELD_TYPES
44def is_core_schema_field(
45 schema: CoreSchemaOrField,
46) -> TypeGuard[CoreSchemaField]:
47 return schema['type'] in _CORE_SCHEMA_FIELD_TYPES
50def is_function_with_inner_schema(
51 schema: CoreSchemaOrField,
52) -> TypeGuard[FunctionSchemaWithInnerSchema]:
53 return schema['type'] in _FUNCTION_WITH_INNER_SCHEMA_TYPES
56def as_ser_schema(schema: CoreSchema) -> core_schema.SerSchema:
57 """Return a schema suitable for the `'serialization'` key of a core schema, from an arbitrary core schema.
59 Any core schema can be used as a serialization schema, except `'function-plain'` and `'function-wrap'`
60 schemas, as these types are already used by the plain and wrap serializer *function* schemas.
61 For these, mimic what pydantic-core would do if they were used as the *main* schema:
63 - if the function schema has a `'serialization'` schema, use it.
64 - otherwise, a `'function-plain'` schema is serialized as `'any'`, and a `'function-wrap'` schema
65 is serialized using its inner schema.
67 Note that `schema` is expected to be a *core* schema (e.g. as returned by the schema generator),
68 not an existing `'serialization'` schema.
69 """
70 if schema['type'] == 'function-plain' or schema['type'] == 'function-wrap':
71 if (ser_schema := schema.get('serialization')) is not None:
72 return ser_schema
73 if schema['type'] == 'function-plain':
74 return core_schema.any_schema()
75 return as_ser_schema(schema['schema'])
76 return schema
79def is_list_like_schema_with_items_schema(
80 schema: CoreSchema,
81) -> TypeGuard[core_schema.ListSchema | core_schema.SetSchema | core_schema.FrozenSetSchema]:
82 return schema['type'] in _LIST_LIKE_SCHEMA_WITH_ITEMS_TYPES
85def get_ref(s: core_schema.CoreSchema) -> None | str:
86 """Get the ref from the schema if it has one.
87 This exists just for type checking to work correctly.
88 """
89 return s.get('ref', None)
92def _clean_schema_for_pretty_print(obj: Any, strip_metadata: bool = True) -> Any: # pragma: no cover
93 """A utility function to remove irrelevant information from a core schema."""
94 if isinstance(obj, Mapping):
95 new_dct = {}
96 for k, v in obj.items():
97 if k == 'metadata' and strip_metadata:
98 new_metadata = {}
100 for meta_k, meta_v in v.items():
101 if meta_k in ('pydantic_js_functions', 'pydantic_js_annotation_functions'):
102 new_metadata['js_metadata'] = '<stripped>'
103 else:
104 new_metadata[meta_k] = _clean_schema_for_pretty_print(meta_v, strip_metadata=strip_metadata)
106 if list(new_metadata.keys()) == ['js_metadata']:
107 new_metadata = {'<stripped>'}
109 new_dct[k] = new_metadata
110 # Remove some defaults:
111 elif k in ('custom_init', 'root_model') and not v:
112 continue
113 else:
114 new_dct[k] = _clean_schema_for_pretty_print(v, strip_metadata=strip_metadata)
116 return new_dct
117 elif isinstance(obj, Sequence) and not isinstance(obj, str):
118 return [_clean_schema_for_pretty_print(v, strip_metadata=strip_metadata) for v in obj]
119 else:
120 return obj
123def pretty_print_core_schema(
124 val: Any,
125 *,
126 console: Console | None = None,
127 max_depth: int | None = None,
128 strip_metadata: bool = True,
129) -> None: # pragma: no cover
130 """Pretty-print a core schema using the `rich` library.
132 Args:
133 val: The core schema to print, or a Pydantic model/dataclass/type adapter
134 (in which case the cached core schema is fetched and printed).
135 console: A rich console to use when printing. Defaults to the global rich console instance.
136 max_depth: The number of nesting levels which may be printed.
137 strip_metadata: Whether to strip metadata in the output. If `True` any known core metadata
138 attributes will be stripped (but custom attributes are kept). Defaults to `True`.
139 """
140 # lazy import:
141 from rich.pretty import pprint
143 # circ. imports:
144 from pydantic import BaseModel, TypeAdapter
145 from pydantic.dataclasses import is_pydantic_dataclass
147 if (inspect.isclass(val) and issubclass(val, BaseModel)) or is_pydantic_dataclass(val):
148 val = val.__pydantic_core_schema__
149 if isinstance(val, TypeAdapter):
150 val = val.core_schema
151 cleaned_schema = _clean_schema_for_pretty_print(val, strip_metadata=strip_metadata)
153 pprint(cleaned_schema, console=console, max_depth=max_depth)
156pps = pretty_print_core_schema