22 lines
632 B
Python
22 lines
632 B
Python
from contextlib import contextmanager
|
|
import signal
|
|
def raise_timeout(signum, frame):
|
|
raise TimeoutError
|
|
|
|
@contextmanager
|
|
def timeout(time:int, exception = TimeoutError):
|
|
# Register a function to raise a TimeoutError on the signal.
|
|
signal.signal(signal.SIGALRM, raise_timeout)
|
|
# Schedule the signal to be sent after ``time``.
|
|
signal.alarm(time)
|
|
|
|
try:
|
|
yield
|
|
except TimeoutError:
|
|
print('TIMED OUT')
|
|
raise exception
|
|
finally:
|
|
# Unregister the signal so it won't be triggered
|
|
# if the timeout is not reached.
|
|
signal.signal(signal.SIGALRM, signal.SIG_IGN)
|