Initial commit

venv/lib/python3.10/site-packages/stamina/instrumentation/_hooks.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from typing import Iterable
+
+from ._data import RetryHook, RetryHookFactory
+from ._logging import LoggingOnRetryHook
+from ._prometheus import PrometheusOnRetryHook
+from ._structlog import StructlogOnRetryHook
+
+
+def init_hooks(
+    maybe_delayed: tuple[RetryHook | RetryHookFactory, ...],
+) -> tuple[RetryHook, ...]:
+    """
+    Execute delayed hook factories and return a tuple of finalized hooks.
+    """
+    hooks = []
+    for hook_or_factory in maybe_delayed:
+        if isinstance(hook_or_factory, RetryHookFactory):
+            hooks.append(hook_or_factory.hook_factory())
+        else:
+            hooks.append(hook_or_factory)
+
+    return tuple(hooks)
+
+
+def get_default_hooks() -> tuple[RetryHookFactory, ...]:
+    """
+    Return the default hooks according to availability.
+    """
+    hooks = []
+
+    try:
+        import prometheus_client  # noqa: F401
+
+        hooks.append(PrometheusOnRetryHook)
+    except ImportError:
+        pass
+
+    try:
+        import structlog  # noqa: F401
+
+        hooks.append(StructlogOnRetryHook)
+    except ImportError:
+        hooks.append(LoggingOnRetryHook)
+
+    return tuple(hooks)
+
+
+def set_on_retry_hooks(
+    hooks: Iterable[RetryHook | RetryHookFactory] | None,
+) -> None:
+    """
+    Set hooks that are called after a retry has been scheduled.
+
+    Args:
+        hooks:
+            Hooks to call after a retry has been scheduled. Passing None resets
+            to default. To deactivate instrumentation, pass an empty iterable.
+
+    .. versionadded:: 23.2.0
+    """
+    from .._config import CONFIG
+
+    CONFIG.on_retry = tuple(hooks) if hooks is not None else hooks  # type: ignore[assignment,arg-type]
+
+
+def get_on_retry_hooks() -> tuple[RetryHook, ...]:
+    """
+    Get hooks that are called after a retry has been scheduled.
+
+    Returns:
+        Hooks that will run if a retry is scheduled. Factories are called if
+        they haven't already.
+
+    .. versionadded:: 23.2.0
+    """
+    from .._config import CONFIG
+
+    return CONFIG.on_retry