90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
from functools import wraps
|
|
from typing import Any, Callable, Dict, TypeVar, cast
|
|
|
|
from typeguard import TypeCheckError, check_type
|
|
|
|
from ..errors import ArgumentValidationError
|
|
from ..helper import get_arguments, get_function_metadata_store
|
|
from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
|
|
from ..tracing.tracing_context import TracingContext
|
|
|
|
F = TypeVar("F", bound=Callable)
|
|
|
|
|
|
def parameter(
|
|
parameter_name: str,
|
|
*,
|
|
validate: Callable[[Any], bool] = lambda _: True,
|
|
disable_logging: bool = False,
|
|
) -> Callable[[F], F]:
|
|
"""Control the validation and logging of function parameters.
|
|
|
|
Basically, a parameter decorator. Unfortunately, Python does not have that concept,
|
|
thus, it's a method decorator that expects the name of the to-be-decorated
|
|
parameter.
|
|
|
|
Examples:
|
|
>>> @parameter('a')
|
|
... def my_function(a: int):
|
|
... return a + 2
|
|
>>> my_function(4)
|
|
6
|
|
>>> my_function('3')
|
|
Traceback (most recent call last):
|
|
...
|
|
TypeError: argument a is not of the expected type: ...
|
|
|
|
>>> @parameter('positive_a', validate=lambda v: v > 0)
|
|
... def my_function(positive_a: int):
|
|
... return a + 2
|
|
>>> my_function(-1)
|
|
Traceback (most recent call last):
|
|
...
|
|
great_ai.errors.argument_validation_error.ArgumentValidationError: ...
|
|
|
|
Args:
|
|
parameter_name: Name of parameter to consider.
|
|
validate: Optional validate to run against the concrete argument.
|
|
ArgumentValidationError is thrown when the return value is False.
|
|
disable_logging: Do not save the value in any active TracingContext.
|
|
Returns:
|
|
A decorator for argument validation.
|
|
"""
|
|
|
|
def decorator(func: F) -> F:
|
|
get_function_metadata_store(func).input_parameter_names.append(parameter_name)
|
|
assert_function_is_not_finalised(func)
|
|
|
|
actual_name = f"arg:{parameter_name}"
|
|
|
|
@wraps(func)
|
|
def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
|
|
arguments = get_arguments(func, args, kwargs)
|
|
argument = arguments.get(parameter_name)
|
|
|
|
expected_type = func.__annotations__.get(parameter_name)
|
|
|
|
if expected_type is not None:
|
|
try:
|
|
check_type(argument, expected_type)
|
|
except TypeCheckError as error:
|
|
raise TypeError(
|
|
f"argument {parameter_name} is not of the expected type: {error}"
|
|
) from error
|
|
|
|
if not validate(argument):
|
|
raise ArgumentValidationError(
|
|
f"Argument {parameter_name} in {func.__name__} did not pass validation"
|
|
)
|
|
|
|
context = TracingContext.get_current_tracing_context()
|
|
if context and not disable_logging:
|
|
context.log_value(name=f"{actual_name}:value", value=argument)
|
|
if isinstance(argument, str):
|
|
context.log_value(name=f"{actual_name}:length", value=len(argument))
|
|
|
|
return func(*args, **kwargs)
|
|
|
|
return cast(F, wrapper)
|
|
|
|
return decorator
|