027133b3创建于 5月18日历史提交
from __future__ import annotations

from contextlib import contextmanager
from threading import Lock
from typing import ClassVar


_instance_id = 0
_lock = Lock()


def _make_unique(name: str):
    global _instance_id  # pylint: disable=global-statement
    with _lock:  # Guard against race conditions
        new_name = f"{name}__{_instance_id}"
        _instance_id += 1
        return new_name


class _ErrorLimiterTrigger(BaseException):  # Use BaseException to avoid being caught by "except Exception:".
    def __init__(self, name: str, *args):
        super().__init__(*args)
        self.name = name.rsplit("__", 1)[0]


class _ErrorLimiter:
    _store: ClassVar[dict[str, int]] = {}

    @classmethod
    def start(cls, name: str, limit: int = 5):
        cls._store[name] = limit

    @classmethod
    def notify(cls, key: str):
        cls._store[key] = cls._store[key] - 1
        if cls._store[key] <= 0:
            raise _ErrorLimiterTrigger(key)

    @classmethod
    def end(cls, name: str):
        cls._store.pop(name)


class ErrorLimiterAbort(RuntimeError):
    def __init__(self, msg: str):
        super().__init__(msg)


@contextmanager
def limit_errors(name: str, limit: int = 5):
    """Limiter for aborting execution after being triggered a specified number of times (default 5).

    >>> with limit_errors("identifier", limit=5) as elimit:
    >>>     while do_thing():
    >>>         if (something_bad):
    >>>             print("Something bad happened")
    >>>             elimit()  # In this example, raises ErrorLimiterAbort on the 5th call
    >>>         try:
    >>>             something_broken()
    >>>         except Exception:
    >>>             print("Encountered an exception")
    >>>             elimit()  # Count is shared across all calls

    Args:
        name (str): Identifier.
        limit (int, optional): Abort after `limit` number of triggers. Defaults to 5.

    Raises:
        ErrorLimiterAbort: Subclass of RuntimeException.

    Yields:
        Callable: Notification function to indicate that an error occurred.
    """
    name_id = _make_unique(name)
    try:
        _ErrorLimiter.start(name_id, limit)
        yield lambda: _ErrorLimiter.notify(name_id)
    except _ErrorLimiterTrigger as e:
        raise ErrorLimiterAbort(f"HALTING. Too many errors during '{e.name}'") from None
    finally:
        _ErrorLimiter.end(name_id)