diff --git a/tests/test_delayed_signals.py b/tests/test_delayed_signals.py new file mode 100644 index 00000000..d840bb0a --- /dev/null +++ b/tests/test_delayed_signals.py @@ -0,0 +1,53 @@ +import os +from signal import signal, SIGHUP, SIGINT +import time +from types import FrameType +import unittest + +from zipline.utils.delayed_signals import delayed_signals + +class DelayedSignals(unittest.TestCase): + def handler(self, signum, frame): + print "Got signal " + str(signum) + self.got[signum] = time.time() + self.assertTrue(isinstance(frame, FrameType)) + + def setUp(self): + signal(SIGHUP, self.handler) + signal(SIGINT, self.handler) + + def reset(self): + self.got = {} + + def test_delayed_signals(self): + self.reset() + with delayed_signals([SIGHUP]): + os.kill(os.getpid(), SIGHUP) + time.sleep(2) + self.assertTrue(self.got[SIGHUP]) + self.assertTrue(time.time() - self.got[SIGHUP] < 2) + + def test_immediate_signals(self): + self.reset() + os.kill(os.getpid(), SIGHUP) + time.sleep(2) + self.assertTrue(self.got[SIGHUP]) + self.assertTrue(time.time() - self.got[SIGHUP] > 1) + + def test_multiple_signals(self): + self.reset() + with delayed_signals([SIGHUP, SIGINT]): + os.kill(os.getpid(), SIGINT) + self.assertFalse(SIGHUP in self.got) + self.assertTrue(SIGINT in self.got) + + @delayed_signals([SIGHUP]) + def kill_and_sleep(self): + os.kill(os.getpid(), SIGHUP) + time.sleep(2) + + def test_decorator(self): + self.reset() + self.kill_and_sleep() + self.assertTrue(SIGHUP in self.got) + self.assertTrue(time.time() - self.got[SIGHUP] < 2) diff --git a/zipline/lines.py b/zipline/lines.py index 1c3a558f..d31c47d6 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -172,6 +172,8 @@ class SimulatedTrading(object): def close(self): log.info("Closing Simulation: {id}".format(id=self.sim_id)) + if self.results_socket: + self.results_socket.close() if self.proc and self.send_sighup: ppid = os.getppid() if self.success: diff --git a/zipline/utils/delayed_signals.py b/zipline/utils/delayed_signals.py new file mode 100644 index 00000000..9ca8c811 --- /dev/null +++ b/zipline/utils/delayed_signals.py @@ -0,0 +1,40 @@ +from functools import wraps +from signal import signal + +class delayed_signals(object): + """ + Utility to temporary intercept one or more signals while a function or code + block is executed, restore their signal handlers at the end of execution, + and invoke them if the signals were in fact received during execution. + + Can be used either as a decorator or a context manager. + + Pass in an iterable of signals to intercept. + """ + + def handler(self, signum, frame=None): + self.got.append({'signum': signum, 'frame': frame}) + + def __init__(self, signals): + self.signals = signals + self.handlers = {} + self.got = [] + + def __enter__(self): + for signum in self.signals: + # signal() returns the old signal handler + self.handlers[signum] = signal(signum, self.handler) + + def __exit__(self, time, value, traceback): + for signum, handler in self.handlers.items(): + signal(signum, handler) + for signum, frame in ((i['signum'], i['frame']) for i in self.got): + self.handlers[signum](signum, frame) + + def __call__(self, fn): + @wraps(fn) + def call_fn(*args, **kwargs): + with self: + outval = fn(*args, **kwargs) + return outval + return call_fn