_schema_validator.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. """Pluggable schema validator for pydantic."""
  2. from __future__ import annotations
  3. import functools
  4. from collections.abc import Iterable
  5. from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar
  6. from pydantic_core import CoreConfig, CoreSchema, SchemaValidator, ValidationError
  7. from typing_extensions import ParamSpec
  8. if TYPE_CHECKING:
  9. from . import BaseValidateHandlerProtocol, PydanticPluginProtocol, SchemaKind, SchemaTypePath
  10. P = ParamSpec('P')
  11. R = TypeVar('R')
  12. Event = Literal['on_validate_python', 'on_validate_json', 'on_validate_strings']
  13. events: list[Event] = list(Event.__args__) # type: ignore
  14. def create_schema_validator(
  15. schema: CoreSchema,
  16. schema_type: Any,
  17. schema_type_module: str,
  18. schema_type_name: str,
  19. schema_kind: SchemaKind,
  20. config: CoreConfig | None = None,
  21. plugin_settings: dict[str, Any] | None = None,
  22. _use_prebuilt: bool = True,
  23. ) -> SchemaValidator | PluggableSchemaValidator:
  24. """Create a `SchemaValidator` or `PluggableSchemaValidator` if plugins are installed.
  25. Returns:
  26. If plugins are installed then return `PluggableSchemaValidator`, otherwise return `SchemaValidator`.
  27. """
  28. from . import SchemaTypePath
  29. from ._loader import get_plugins
  30. plugins = get_plugins()
  31. if plugins:
  32. return PluggableSchemaValidator(
  33. schema,
  34. schema_type,
  35. SchemaTypePath(schema_type_module, schema_type_name),
  36. schema_kind,
  37. config,
  38. plugins,
  39. plugin_settings or {},
  40. _use_prebuilt=_use_prebuilt,
  41. )
  42. else:
  43. return SchemaValidator(schema, config, _use_prebuilt=_use_prebuilt)
  44. class PluggableSchemaValidator:
  45. """Pluggable schema validator."""
  46. __slots__ = '_schema_validator', 'validate_json', 'validate_python', 'validate_strings'
  47. def __init__(
  48. self,
  49. schema: CoreSchema,
  50. schema_type: Any,
  51. schema_type_path: SchemaTypePath,
  52. schema_kind: SchemaKind,
  53. config: CoreConfig | None,
  54. plugins: Iterable[PydanticPluginProtocol],
  55. plugin_settings: dict[str, Any],
  56. _use_prebuilt: bool = True,
  57. ) -> None:
  58. self._schema_validator = SchemaValidator(schema, config, _use_prebuilt=_use_prebuilt)
  59. python_event_handlers: list[BaseValidateHandlerProtocol] = []
  60. json_event_handlers: list[BaseValidateHandlerProtocol] = []
  61. strings_event_handlers: list[BaseValidateHandlerProtocol] = []
  62. for plugin in plugins:
  63. try:
  64. p, j, s = plugin.new_schema_validator(
  65. schema, schema_type, schema_type_path, schema_kind, config, plugin_settings
  66. )
  67. except TypeError as e: # pragma: no cover
  68. raise TypeError(f'Error using plugin `{plugin.__module__}:{plugin.__class__.__name__}`: {e}') from e
  69. if p is not None:
  70. python_event_handlers.append(p)
  71. if j is not None:
  72. json_event_handlers.append(j)
  73. if s is not None:
  74. strings_event_handlers.append(s)
  75. self.validate_python = build_wrapper(self._schema_validator.validate_python, python_event_handlers)
  76. self.validate_json = build_wrapper(self._schema_validator.validate_json, json_event_handlers)
  77. self.validate_strings = build_wrapper(self._schema_validator.validate_strings, strings_event_handlers)
  78. def __getattr__(self, name: str) -> Any:
  79. return getattr(self._schema_validator, name)
  80. def build_wrapper(func: Callable[P, R], event_handlers: list[BaseValidateHandlerProtocol]) -> Callable[P, R]:
  81. if not event_handlers:
  82. return func
  83. else:
  84. on_enters = tuple(h.on_enter for h in event_handlers if filter_handlers(h, 'on_enter'))
  85. on_successes = tuple(h.on_success for h in event_handlers if filter_handlers(h, 'on_success'))
  86. on_errors = tuple(h.on_error for h in event_handlers if filter_handlers(h, 'on_error'))
  87. on_exceptions = tuple(h.on_exception for h in event_handlers if filter_handlers(h, 'on_exception'))
  88. @functools.wraps(func)
  89. def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
  90. for on_enter_handler in on_enters:
  91. on_enter_handler(*args, **kwargs)
  92. try:
  93. result = func(*args, **kwargs)
  94. except ValidationError as error:
  95. for on_error_handler in on_errors:
  96. on_error_handler(error)
  97. raise
  98. except Exception as exception:
  99. for on_exception_handler in on_exceptions:
  100. on_exception_handler(exception)
  101. raise
  102. else:
  103. for on_success_handler in on_successes:
  104. on_success_handler(result)
  105. return result
  106. return wrapper
  107. def filter_handlers(handler_cls: BaseValidateHandlerProtocol, method_name: str) -> bool:
  108. """Filter out handler methods which are not implemented by the plugin directly - e.g. those that are missing
  109. or are inherited from the protocol.
  110. """
  111. handler = getattr(handler_cls, method_name, None)
  112. if handler is None:
  113. return False
  114. elif handler.__module__ == 'pydantic.plugin':
  115. # this is the original handler, from the protocol due to runtime inheritance
  116. # we don't want to call it
  117. return False
  118. else:
  119. return True