Daehyun
12/8/2017 - 4:22 PM

Timeout

from signal import signal, SIGALRM, alarm

from cytoolz import curry

class Timeout:
    def __init__(self, sec):
        self.__sec = sec

    def __enter__(self):
        signal(SIGALRM, self.raise_error)
        alarm(self.__sec)

    def __exit__(self, *args):
        alarm(0)

    @staticmethod
    def raise_error(*args):
        raise TimeoutError


@curry
def timeout(sec, func):
    def curried(*args, **kwargs):
        signal(SIGALRM, Timeout.raise_error)
        alarm(sec)
        ret = func(*args, **kwargs)
        alarm(0)
        return ret
    return curried