196 lines
4.7 KiB
Python
196 lines
4.7 KiB
Python
# SPDX-FileCopyrightText: 2022 Hynek Schlawack <hs@ox.cx>
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
from __future__ import annotations
|
|
|
|
from threading import Lock
|
|
from types import TracebackType
|
|
from typing import Callable
|
|
|
|
from .instrumentation import RetryHookFactory
|
|
from .instrumentation._hooks import get_default_hooks, init_hooks
|
|
from .typing import RetryHook
|
|
|
|
|
|
class _Testing:
|
|
"""
|
|
Test mode specification.
|
|
|
|
Strictly private.
|
|
"""
|
|
|
|
__slots__ = ("attempts", "cap")
|
|
|
|
attempts: int
|
|
cap: bool
|
|
|
|
def __init__(self, attempts: int, cap: bool) -> None:
|
|
self.attempts = attempts
|
|
self.cap = cap
|
|
|
|
def get_attempts(self, non_testing_attempts: int | None) -> int:
|
|
"""
|
|
Get the number of attempts to use.
|
|
|
|
Args:
|
|
non_testing_attempts: The number of attempts specified by the user.
|
|
|
|
Returns:
|
|
The number of attempts to use.
|
|
"""
|
|
if self.cap:
|
|
return min(self.attempts, non_testing_attempts or self.attempts)
|
|
|
|
return self.attempts
|
|
|
|
|
|
class _Config:
|
|
"""
|
|
Global stamina configuration.
|
|
|
|
Strictly private.
|
|
"""
|
|
|
|
__slots__ = (
|
|
"_get_on_retry",
|
|
"_is_active",
|
|
"_on_retry",
|
|
"_testing",
|
|
"lock",
|
|
)
|
|
|
|
lock: Lock
|
|
_is_active: bool
|
|
_testing: _Testing | None
|
|
_on_retry: (
|
|
tuple[RetryHook, ...] | tuple[RetryHook | RetryHookFactory, ...] | None
|
|
)
|
|
_get_on_retry: Callable[[], tuple[RetryHook, ...]]
|
|
|
|
def __init__(self, lock: Lock) -> None:
|
|
self.lock = lock
|
|
self._is_active = True
|
|
self._testing = None
|
|
|
|
# Prepare delayed initialization.
|
|
self._on_retry = None
|
|
self._get_on_retry = self._init_on_first_retry
|
|
|
|
@property
|
|
def is_active(self) -> bool:
|
|
return self._is_active
|
|
|
|
@is_active.setter
|
|
def is_active(self, value: bool) -> None:
|
|
with self.lock:
|
|
self._is_active = value
|
|
|
|
@property
|
|
def testing(self) -> _Testing | None:
|
|
return self._testing
|
|
|
|
@testing.setter
|
|
def testing(self, value: _Testing | None) -> None:
|
|
with self.lock:
|
|
self._testing = value
|
|
|
|
@property
|
|
def on_retry(self) -> tuple[RetryHook, ...]:
|
|
return self._get_on_retry()
|
|
|
|
@on_retry.setter
|
|
def on_retry(
|
|
self, value: tuple[RetryHook | RetryHookFactory, ...] | None
|
|
) -> None:
|
|
with self.lock:
|
|
self._get_on_retry = self._init_on_first_retry
|
|
self._on_retry = value
|
|
|
|
def _init_on_first_retry(self) -> tuple[RetryHook, ...]:
|
|
"""
|
|
Perform delayed initialization of on_retry hooks.
|
|
"""
|
|
with self.lock:
|
|
# Ensure hooks didn't init while waiting for the lock.
|
|
if self._get_on_retry == self._init_on_first_retry:
|
|
if self._on_retry is None:
|
|
self._on_retry = get_default_hooks()
|
|
|
|
self._on_retry = init_hooks(self._on_retry)
|
|
|
|
self._get_on_retry = lambda: self._on_retry # type: ignore[assignment, return-value]
|
|
|
|
return self._on_retry # type: ignore[return-value]
|
|
|
|
|
|
CONFIG = _Config(Lock())
|
|
|
|
|
|
def is_active() -> bool:
|
|
"""
|
|
Check whether retrying is active.
|
|
|
|
Returns:
|
|
Whether retrying is active.
|
|
"""
|
|
return CONFIG.is_active
|
|
|
|
|
|
def set_active(active: bool) -> None:
|
|
"""
|
|
Activate or deactivate retrying.
|
|
|
|
Is idempotent and can be called repeatedly with the same value.
|
|
"""
|
|
CONFIG.is_active = bool(active)
|
|
|
|
|
|
def is_testing() -> bool:
|
|
"""
|
|
Check whether test mode is enabled.
|
|
|
|
.. versionadded:: 24.3.0
|
|
"""
|
|
return CONFIG.testing is not None
|
|
|
|
|
|
class _RestoreTestingCM:
|
|
def __init__(self, old: _Testing | None) -> None:
|
|
self.old = old
|
|
|
|
def __enter__(self) -> None:
|
|
pass
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
) -> None:
|
|
CONFIG.testing = self.old
|
|
|
|
|
|
def set_testing(
|
|
testing: bool, *, attempts: int = 1, cap: bool = False
|
|
) -> _RestoreTestingCM:
|
|
"""
|
|
Activate or deactivate test mode.
|
|
|
|
In testing mode, backoffs are disabled, and attempts are set to *attempts*.
|
|
|
|
If *cap* is True, the number of attempts is not set but capped at
|
|
*attempts*. This means that if *attempts* is greater than the number of
|
|
attempts specified by the user, the user's value is used.
|
|
|
|
Is idempotent and can be called repeatedly with the same values.
|
|
|
|
.. versionadded:: 24.3.0
|
|
.. versionadded:: 25.1.0 *cap*
|
|
.. versionadded:: 25.1.0 Can be used as a context manager.
|
|
"""
|
|
old = CONFIG.testing
|
|
CONFIG.testing = _Testing(attempts, cap) if testing else None
|
|
|
|
return _RestoreTestingCM(old)
|