1 files changed, 65 insertions, 0 deletions
diff --git a/util/thread_timeout.py b/util/thread_timeout.py
new file mode 100644
index 0000000..9a2637c
--- /dev/null
+++ b/util/thread_timeout.py
@@ -0,0 +1,65 @@
+from threading import Thread
+import inspect
+import ctypes
+from functools import wraps
+
+
+def _async_raise(tid, exctype):
+    """raises the exception, performs cleanup if needed"""
+    tid = ctypes.c_long(tid)
+    if not inspect.isclass(exctype):
+        exctype = type(exctype)
+    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
+    if res == 0:
+        raise ValueError("invalid thread id")
+    elif res != 1:
+        # """if it returns a number greater than one, you're in trouble,
+        # and you should call it again with exc=NULL to revert the effect"""
+        ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
+        raise SystemError("PyThreadState_SetAsyncExc failed")
+
+
+def stop_thread(thread):
+    _async_raise(thread.ident, SystemExit)
+
+
+class TimeoutException(Exception):
+    # print("timeout!")
+    pass
+
+
+ThreadStop = stop_thread
+
+
+def time_limited(timeout):
+    def decorator(function):
+        @wraps(function)
+        def wrapped_function(*args, **kwargs):
+            class TimeLimited(Thread):
+                def __init__(self):
+                    Thread.__init__(self)
+                    self.error = None
+                    self.result = None
+
+                def run(self):
+                    self.result = function(*args, **kwargs)
+
+                def stop(self):
+                    if self.is_alive():
+                        ThreadStop(self)
+
+            t = TimeLimited()
+            t.start()
+            t.join(timeout)
+            if isinstance(t.error, TimeoutException):
+                t.stop()
+                raise TimeoutException('timeout for %s' % (repr(function)))
+            if t.is_alive():
+                t.stop()
+                raise TimeoutException('timeout for %s' % (repr(function)))
+            if t.error is None:
+                return t.result
+
+        return wrapped_function
+
+    return decorator