diff --git a/tomorrow/tomorrow.py b/tomorrow/tomorrow.py index bb3e679..e2f3fb4 100644 --- a/tomorrow/tomorrow.py +++ b/tomorrow/tomorrow.py @@ -1,23 +1,46 @@ from functools import wraps +import concurrent.futures from concurrent.futures import ThreadPoolExecutor -class Tomorrow(): +EXECUTORS = [] + +class Tomorrow(): + """ + Result wrapper that captures __getattr__ calls in order to + wait on futures when any attribute of the Tomorrow object is accessed + """ def __init__(self, future, timeout): self._future = future self._timeout = timeout def __getattr__(self, name): - result = self._wait() - return result.__getattribute__(name) + """ + When an attribute of the Tomorrow object is accessed, resolve obj._future + """ + try: + result = self._wait() + return result.__getattribute__(name) + except KeyboardInterrupt: + exit() + raise - def _wait(self): - return self._future.result(self._timeout) + def _wait(self, timeout=None): + """ + Wait on a future that only causes side effects and + has no return value that is accessed with a __getattr__ call + """ + if not timeout: + timeout = self._timeout + return self._future.result(timeout) def async(n, base_type, timeout=None): + """ + Base function for @threads and (eventually) @processes decorators + """ def decorator(f): if isinstance(n, int): pool = base_type(n) @@ -28,6 +51,9 @@ def decorator(f): "Invalid type: %s" % type(base_type) ) + + EXECUTORS.append(pool) + @wraps(f) def wrapped(*args, **kwargs): return Tomorrow( @@ -39,4 +65,18 @@ def wrapped(*args, **kwargs): def threads(n, timeout=None): + """ + Decorator for threaded execution + """ return async(n, ThreadPoolExecutor, timeout) + + +def exit(): + """ + Clean up all future objects and resume execution in the main thread + Allows for program exit when an interrupt signal is received + """ + for executor in EXECUTORS: + executor._threads.clear() + concurrent.futures.thread._threads_queues.clear() + \ No newline at end of file