arguments_schema.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. """Experimental module exposing a function to generate a core schema that validates callable arguments."""
  2. from __future__ import annotations
  3. from collections.abc import Callable
  4. from typing import Any, Literal
  5. from pydantic_core import CoreSchema
  6. from pydantic import ConfigDict
  7. from pydantic._internal import _config, _generate_schema, _namespace_utils
  8. def generate_arguments_schema(
  9. func: Callable[..., Any],
  10. schema_type: Literal['arguments', 'arguments-v3'] = 'arguments-v3',
  11. parameters_callback: Callable[[int, str, Any], Literal['skip'] | None] | None = None,
  12. config: ConfigDict | None = None,
  13. ) -> CoreSchema:
  14. """Generate the schema for the arguments of a function.
  15. Args:
  16. func: The function to generate the schema for.
  17. schema_type: The type of schema to generate.
  18. parameters_callback: A callable that will be invoked for each parameter. The callback
  19. should take three required arguments: the index, the name and the type annotation
  20. (or [`Parameter.empty`][inspect.Parameter.empty] if not annotated) of the parameter.
  21. The callback can optionally return `'skip'`, so that the parameter gets excluded
  22. from the resulting schema.
  23. config: The configuration to use.
  24. Returns:
  25. The generated schema.
  26. """
  27. generate_schema = _generate_schema.GenerateSchema(
  28. _config.ConfigWrapper(config),
  29. ns_resolver=_namespace_utils.NsResolver(namespaces_tuple=_namespace_utils.ns_for_function(func)),
  30. )
  31. if schema_type == 'arguments':
  32. schema = generate_schema._arguments_schema(func, parameters_callback) # pyright: ignore[reportArgumentType]
  33. else:
  34. schema = generate_schema._arguments_v3_schema(func, parameters_callback) # pyright: ignore[reportArgumentType]
  35. return generate_schema.clean_schema(schema)