Source code for hfai.utils.time

import signal
from functools import wraps


class TimeoutError(Exception):

    def __init__(self, f, sec) -> None:
        msg = f"{f.__name__} is timeout for {sec} seconds"
        self.f = f
        self.sec = sec
        super().__init__(msg)


def timeout(seconds=0):
    """
    给函数设置超时限制

    Args:
        seconds (int): 超时的秒数

    Examples:

    .. code-block:: python

        import time, hfai
        from hfai.utils import TimeoutError

        @hfai.utils.timeout(10)
        def print_hello():
            while True:
                print("hello")
                time.sleep(1)

        try:
            print_hello()
        except TimeoutError as e:
            print(e)

    """
    assert isinstance(seconds, int) and seconds >= 0

    def timeout_decorator(f):
        def timeout_handler(*unsed):
            raise TimeoutError(f, seconds)

        @wraps(f)
        def new_f(*args, **kwargs):
            if seconds == 0:
                return f(*args, **kwargs)

            # set the timeout handler
            signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(seconds)

            try:
                result = f(*args, **kwargs)
            finally:
                signal.alarm(0)

            return result

        return new_f

    return timeout_decorator