Source code for core_mixins.decorators.timeout_signal

# -*- coding: utf-8 -*-

"""
This decorator (alternative to `with_timeout`) ensures a function
does not run for more than a certain amount of seconds
before an error (`TimeoutError`) is raised.
"""

import functools
import platform
import signal
import threading
import typing




def with_timeout_signal(
    fcn: typing.Optional[typing.Callable] = None,
    timeout: float = 10,
) -> typing.Callable:
    """
    It executes a function that will time out after the
    specified amount of seconds using `signal`. Could be an
    alternative for `with_timeout` because does not require
    spanning another process or the use of a queue.

    Caveats:
      - Only works on Unix/Linux systems (SIGALRM not available on Windows).
      - Must be called from the main thread (signal.alarm limitation).
      - Not re-entrant: nested or concurrent decorated functions will interfere.
      - Not compatible with async functions or async contexts.
      - May conflict with other code using SIGALRM.

    :param fcn: The function being decorated.
    :param timeout: The seconds before time out.
    :return: The wrapped function.

    :raises TimeoutError: If the decorated function exceeds the timeout duration.
    :raises RuntimeError: If used on Windows or from a non-main thread.
    """

    if platform.system() == "Windows":
        raise RuntimeError(
            "The `with_timeout_signal` decorator is not supported "
            "on Windows. Use `with_timeout` decorator instead."
        )

    def decorator_timeout(_fcn: typing.Callable):
        _fcn_name = getattr(_fcn, "__name__", repr(_fcn))

        def timeout_handler(_signum, _frame):
            raise TimeoutError(
                f"Function '{_fcn_name}' timed out "
                f"after {timeout} seconds."
            )

        @functools.wraps(_fcn)
        def wrapper(*args, **kwargs):
            if threading.current_thread() != threading.main_thread():
                raise RuntimeError(
                    f"`with_timeout_signal` can only be used from the "
                    f"main thread. Function '{_fcn_name}' was called "
                    f"from thread: '{threading.current_thread().name}'."
                )

            signal.signal(signal.SIGALRM, timeout_handler)
            signal.setitimer(signal.ITIMER_REAL, timeout)

            try:
                return _fcn(*args, **kwargs)

            finally:
                signal.setitimer(signal.ITIMER_REAL, 0)

        return wrapper

    if not fcn:
        return decorator_timeout

    return decorator_timeout(fcn)