Initial commit

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

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from ._data import RetryDetails, RetryHook, RetryHookFactory
+from ._hooks import get_on_retry_hooks, set_on_retry_hooks
+from ._logging import LoggingOnRetryHook
+from ._prometheus import PrometheusOnRetryHook, get_prometheus_counter
+from ._structlog import StructlogOnRetryHook
+
+
+__all__ = [
+    "LoggingOnRetryHook",
+    "PrometheusOnRetryHook",
+    "RetryDetails",
+    "RetryHook",
+    "RetryHookFactory",
+    "StructlogOnRetryHook",
+    "get_on_retry_hooks",
+    "get_prometheus_counter",
+    "set_on_retry_hooks",
+]

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

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from typing import Callable, Protocol
+
+
+def guess_name(obj: object) -> str:
+    name = getattr(obj, "__qualname__", None) or "<unnamed object>"
+    mod = getattr(obj, "__module__", None) or "<unknown module>"
+
+    if mod == "builtins":
+        return name
+
+    return f"{mod}.{name}"
+
+
+@dataclass(frozen=True)
+class RetryDetails:
+    r"""
+    Details about a retry attempt that are passed into :class:`RetryHook`\ s.
+
+    Attributes:
+        name: Name of the callable that is being retried.
+
+        args: Positional arguments that were passed to the callable.
+
+        kwargs: Keyword arguments that were passed to the callable.
+
+        retry_num:
+            Number of the retry attempt. Starts at 1 after the first failure.
+
+        wait_for:
+            Time in seconds that *stamina* will wait before the next attempt.
+
+        waited_so_far:
+            Time in seconds that *stamina* has waited so far for the current
+            callable.
+
+        caused_by: Exception that caused the retry attempt.
+
+    .. versionadded:: 23.2.0
+    """
+
+    __slots__ = (
+        "args",
+        "caused_by",
+        "kwargs",
+        "name",
+        "retry_num",
+        "wait_for",
+        "waited_so_far",
+    )
+
+    name: str
+    args: tuple[object, ...]
+    kwargs: dict[str, object]
+    retry_num: int
+    wait_for: float
+    waited_so_far: float
+    caused_by: Exception
+
+
+class RetryHook(Protocol):
+    """
+    A callable that gets called after an attempt has failed and a retry has
+    been scheduled.
+
+    This is a :class:`typing.Protocol` that can be implemented by any callable
+    that takes one argument of type :class:`RetryDetails` and returns None.
+
+    If the hook returns a context manager, it will be entered when the retry is
+    scheduled and exited right before the retry is attempted.
+
+    .. versionadded:: 23.2.0
+
+    .. versionadded:: 25.1.0
+       Added support for context managers.
+    """
+
+    def __call__(
+        self, details: RetryDetails
+    ) -> None | AbstractContextManager[None]: ...
+
+
+@dataclass(frozen=True)
+class RetryHookFactory:
+    """
+    Wraps a callable that returns a :class:`RetryHook`.
+
+    They are called on the first scheduled retry and can be used to delay
+    initialization. If you need to pass arguments, you can do that using
+    :func:`functools.partial`.
+
+    .. versionadded:: 23.2.0
+    """
+
+    hook_factory: Callable[[], RetryHook]

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

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

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from ._data import RetryDetails, RetryHook, RetryHookFactory
+
+
+def init_logging(log_level: int = 30) -> RetryHook:
+    """
+    Initialize logging using the standard library.
+
+    Returned hook logs scheduled retries at *log_level*.
+
+    .. versionadded:: 23.2.0
+    """
+    import logging
+
+    logger = logging.getLogger("stamina")
+
+    def log_retries(details: RetryDetails) -> None:
+        logger.log(
+            log_level,
+            "stamina.retry_scheduled",
+            extra={
+                "stamina.callable": details.name,
+                "stamina.args": tuple(repr(a) for a in details.args),
+                "stamina.kwargs": dict(details.kwargs.items()),
+                "stamina.retry_num": details.retry_num,
+                "stamina.caused_by": repr(details.caused_by),
+                "stamina.wait_for": round(details.wait_for, 2),
+                "stamina.waited_so_far": round(details.waited_so_far, 2),
+            },
+        )
+
+    return log_retries
+
+
+LoggingOnRetryHook = RetryHookFactory(init_logging)

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

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ._data import RetryDetails, RetryHook, RetryHookFactory, guess_name
+
+
+if TYPE_CHECKING:
+    from prometheus_client import Counter
+
+
+RETRIES_TOTAL = None
+
+
+def init_prometheus() -> RetryHook:
+    """
+    Initialize Prometheus instrumentation.
+    """
+    from prometheus_client import Counter
+
+    global RETRIES_TOTAL  # noqa: PLW0603
+
+    # Mostly for testing so we can call init_prometheus more than once.
+    if RETRIES_TOTAL is None:
+        RETRIES_TOTAL = Counter(
+            "stamina_retries_total",
+            "Total number of retries.",
+            ("callable", "retry_num", "error_type"),
+        )
+
+    def count_retries(details: RetryDetails) -> None:
+        """
+        Count and log retries for callable *name*.
+        """
+        RETRIES_TOTAL.labels(
+            callable=details.name,
+            retry_num=details.retry_num,
+            error_type=guess_name(details.caused_by.__class__),
+        ).inc()
+
+    return count_retries
+
+
+def get_prometheus_counter() -> Counter | None:
+    """
+    Return the Prometheus counter for the number of retries.
+
+    Returns:
+        If active, the Prometheus `counter
+        <https://github.com/prometheus/client_python>`_ for the number of
+        retries. None otherwise.
+
+    .. versionadded:: 23.2.0
+    """
+    from . import get_on_retry_hooks
+
+    # Finalize the hooks if not done yet.
+    get_on_retry_hooks()
+
+    return RETRIES_TOTAL
+
+
+PrometheusOnRetryHook = RetryHookFactory(init_prometheus)

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

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from ._data import RetryDetails, RetryHook, RetryHookFactory
+
+
+def init_structlog() -> RetryHook:
+    """
+    Initialize structlog instrumentation.
+
+    .. versionadded:: 23.2.0
+    """
+    import structlog
+
+    logger = structlog.get_logger("stamina")
+
+    def log_retries(details: RetryDetails) -> None:
+        logger.warning(
+            "stamina.retry_scheduled",
+            callable=details.name,
+            args=tuple(repr(a) for a in details.args),
+            kwargs=dict(details.kwargs.items()),
+            retry_num=details.retry_num,
+            caused_by=repr(details.caused_by),
+            wait_for=round(details.wait_for, 2),
+            waited_so_far=round(details.waited_so_far, 2),
+        )
+
+    return log_retries
+
+
+StructlogOnRetryHook = RetryHookFactory(init_structlog)