decorator.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. """Decorator utilities."""
  2. import functools
  3. from collections.abc import Callable
  4. from typing import ParamSpec, TypeVar
  5. T = TypeVar("T")
  6. def once(f: Callable[[], T]) -> Callable[[], T]:
  7. """A decorator that calls the function once and caches the result.
  8. Args:
  9. f: The function to call.
  10. Returns:
  11. A function that calls the function once and caches the result.
  12. """
  13. unset = object()
  14. value: object | T = unset
  15. @functools.wraps(f)
  16. def wrapper() -> T:
  17. nonlocal value
  18. value = f() if value is unset else value
  19. return value # pyright: ignore[reportReturnType]
  20. return wrapper
  21. def once_unless_none(f: Callable[[], T | None]) -> Callable[[], T | None]:
  22. """A decorator that calls the function once and caches the result unless it is None.
  23. Args:
  24. f: The function to call.
  25. Returns:
  26. A function that calls the function once and caches the result unless it is None.
  27. """
  28. value: T | None = None
  29. @functools.wraps(f)
  30. def wrapper() -> T | None:
  31. nonlocal value
  32. value = f() if value is None else value
  33. return value
  34. return wrapper
  35. P = ParamSpec("P")
  36. def debug(f: Callable[P, T]) -> Callable[P, T]:
  37. """A decorator that prints the function name, arguments, and result.
  38. Args:
  39. f: The function to call.
  40. Returns:
  41. A function that prints the function name, arguments, and result.
  42. """
  43. @functools.wraps(f)
  44. def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
  45. result = f(*args, **kwargs)
  46. print( # noqa: T201
  47. f"Calling {f.__name__} with args: {args} and kwargs: {kwargs}, result: {result}"
  48. )
  49. return result
  50. return wrapper