From 5a70c7464aaf40ebbfd59eae350840c15f77a868 Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Fri, 17 Aug 2012 11:07:12 -0400
Subject: [PATCH 1/8] update the comments in tradesim

---
 zipline/gens/tradesimulation.py | 39 +++++++++++++++++++++++----------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 22afa34f..13d60194 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -1,3 +1,5 @@
+import signal
+
 from logbook import Logger, Processor
 
 from datetime import datetime, timedelta
@@ -13,10 +15,19 @@ from zipline.gens.utils import hash_args
 
 log = Logger('Trade Simulation')
 
+class AlgoTimeoutException(Exception):
+    def __init__(self):
+        pass
+
+def handle_init_timeout(signum, frame):
+    log.error("Algorithm timed out during initialize.")
+    raise 
+
+
 class TradeSimulationClient(object):
     """
-    Generator that takes the expected output of a merge, a user
-    algorithm, a trading environment, and a simulator style as
+    Generator-style class that takes the expected output of a merge, a
+    user algorithm, a trading environment, and a simulator style as
     arguments.  Pipes the merge stream through a TransactionSimulator
     and a PerformanceTracker, which keep track of the current state of
     our algorithm's simulated universe. Results are fed to the user's
@@ -24,7 +35,7 @@ class TradeSimulationClient(object):
     TransactionSimulator's order book.
 
     TransactionSimulator maintains a dictionary from sids to the
-    unfulfilled orders placed by the user's algorithm.  As trade
+    as-yet unfilled orders placed by the user's algorithm.  As trade
     events arrive, if the algorithm has open orders against the
     trade's sid, the simulator will fill orders up to 25% of market
     cap.  Applied transactions are added to a txn field on the event
@@ -40,9 +51,9 @@ class TradeSimulationClient(object):
     performance report, which is appended to event's perf_report
     field.
 
-    Fully processed events are run through a batcher generator, which
-    batches together events with the same dt field into a single event
-    to be fed to the algo. The portfolio object is repeatedly
+    Fully processed events are fed to AlgorithmSimulator, which
+    batches together events with the same dt field into a single
+    snapshot to be fed to the algo. The portfolio object is repeatedly
     overwritten so that only the most recent snapshot of the universe
     is sent to the algo.
     """
@@ -54,13 +65,14 @@ class TradeSimulationClient(object):
         self.environment = environment
         self.style = sim_style
         self.algo_sim = None
-
+        
         self.warmup_start = self.environment.prior_day_open
         self.algo_start = self.environment.first_open
 
     def get_hash(self):
         """
-        There should only ever be one TSC in the system.
+        There should only ever be one TSC in the system, so
+        we don't bother passing args into the hash.
         """
         return self.__class__.__name__ + hash_args()
 
@@ -92,9 +104,9 @@ class TradeSimulationClient(object):
         with_portfolio = perf_tracker.transform(with_filled_orders)
 
         # Pass the messages from perf along with the trading client's
-        # state into the algorithm for simulation. We provide the
-        # trading client so that the algorithm can place new orders
-        # into the client's order book.
+        # state into the algorithm for simulation. We provide a
+        # pointer to the ordering client's internal state so that the
+        # algorithm can place new orders into the client's order book.
         self.algo_sim = AlgorithmSimulator(
             with_portfolio,
             ordering_client.state,
@@ -273,7 +285,9 @@ class AlgorithmSimulator(object):
 
     def update_current_snapshot(self, event):
         """
-        Update our current snapshot of the universe. Call handle_data if
+        Update our current snapshot of the universe. If event.dt doesn't
+        match our current snapshot's dt, we simulate the current snapshot
+        before processing the event.
         """
         # The new event matches our snapshot dt. Just update the
         # universe and move on.
@@ -326,3 +340,4 @@ class AlgorithmSimulator(object):
         # Update our knowledge of this event's sid
         for field in event.keys():
             self.universe[event.sid][field] = event[field]
+

From a68a48b62e6bb31534a15a29eb17eae65ac0c23c Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sat, 18 Aug 2012 17:07:20 -0400
Subject: [PATCH 2/8] removed deprecated test and refactored tradesimulation
 time compression logic

---
 tests/test_finance.py           |  30 -----
 zipline/gens/tradesimulation.py | 210 +++++++++++++++-----------------
 2 files changed, 98 insertions(+), 142 deletions(-)

diff --git a/tests/test_finance.py b/tests/test_finance.py
index e5f26240..94a98da9 100644
--- a/tests/test_finance.py
+++ b/tests/test_finance.py
@@ -121,36 +121,6 @@ class FinanceTestCase(TestCase):
         zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
         assert_single_position(self, zipline)
 
-    #@timed(DEFAULT_TIMEOUT)
-    def test_sid_filter(self):
-        # Ensure the algorithm's filter prevents events from arriving.
-        # create a test algorithm whose filter will not match any of the
-        # trade events sourced inside the zipline.
-        order_amount = 100
-        order_count = 100
-        no_match_sid = 222
-        test_algo = TestAlgorithm(
-            no_match_sid,
-            order_amount,
-            order_count
-        )
-
-        self.zipline_test_config['trade_count'] = 200
-        self.zipline_test_config['algorithm'] = test_algo
-
-        zipline = SimulatedTrading.create_test_zipline(
-            **self.zipline_test_config
-        )
-        output, transaction_count = drain_zipline(self, zipline)
-
-        #check that the algorithm received no events
-        self.assertEqual(
-            0,
-            transaction_count,
-            "The algorithm should not receive any events due to filtering."
-        )
-
-
     # TODO: write tests for short sales
     # TODO: write a test to do massive buying or shorting.
 
diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 13d60194..6a5e88b2 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -1,11 +1,12 @@
-import signal
 
 from logbook import Logger, Processor
 
 from datetime import datetime, timedelta
 from numbers import Integral
+from itertools import groupby
 
 from zipline import ndict
+from zipline.utils import heartbeat
 
 from zipline.gens.transform import StatefulTransform
 from zipline.finance.trading import TransactionSimulator
@@ -16,13 +17,7 @@ from zipline.gens.utils import hash_args
 log = Logger('Trade Simulation')
 
 class AlgoTimeoutException(Exception):
-    def __init__(self):
-        pass
-
-def handle_init_timeout(signum, frame):
-    log.error("Algorithm timed out during initialize.")
-    raise 
-
+    pass
 
 class TradeSimulationClient(object):
     """
@@ -118,12 +113,17 @@ class TradeSimulationClient(object):
         # calculated by the performance tracker) at the end of each
         # day.  It will also yield a risk report at the end of the
         # simulation.
+
         for message in self.algo_sim:
             yield message
 
 class AlgorithmSimulator(object):
 
-    def __init__(self, stream_in, order_book, algo, algo_start):
+    def __init__(self, 
+                 stream_in, 
+                 order_book, 
+                 algo, 
+                 algo_start):
 
         self.stream_in = stream_in
 
@@ -140,15 +140,15 @@ class AlgorithmSimulator(object):
         self.algo_start = algo_start
 
         # Monkey patch the user algorithm to place orders in the
-        # TransactionSimulator's order book.
+        # TransactionSimulator's order book and use our logger.
         self.algo.set_order(self.order)
-        self.algo.set_logger(Logger("AlgoLog"))
-
+        self.algolog = Logger("AlgoLog")
+        self.algo.set_logger(self.algolog)
 
         # ==============
         # Snapshot Setup
         # ==============
-
+        
         # The algorithm's universe as of our most recent event.
         self.universe = ndict()
 
@@ -159,7 +159,7 @@ class AlgorithmSimulator(object):
         # We don't have a datetime for the current snapshot until we
         # receive a message.
         self.simulation_dt = None
-        self.this_snapshot_dt = None
+        self.snapshot_dt = None
 
         # =============
         # Logging Setup
@@ -168,7 +168,7 @@ class AlgorithmSimulator(object):
         # Processor function for injecting the algo_dt into
         # user prints/logs.
         def inject_algo_dt(record):
-            record.extra['algo_dt'] = self.this_snapshot_dt
+            record.extra['algo_dt'] = self.snapshot_dt
         self.processor = Processor(inject_algo_dt)
 
         # This is a class, which is instantiated later
@@ -225,110 +225,65 @@ class AlgorithmSimulator(object):
         # snapshot time to any log record generated.
 
         with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''):
+            
+            #Set an alarm to go off if initialize takes more than 5 seconds.
+            signal.signal(signal.SIGALRM, self.handle_init_timeout)
+            signal.alarm(5)
             # Call the user's initialize method.
             self.algo.initialize()
+            # Deactivate the alarm.
+            signal.alarm(0)
+            signal.signal(signal.SIGALRM, signal.SIG_DFL)
 
-            for event in self.stream_in:
+            # Group together events with the same dt field. This depends on the
+            # events already being sorted.
+            for date, snapshot in groupby(self.stream_in, lambda e: e.dt):
+                
+                # Set the simulation date to be the first event we see.
+                # This should only occur once, at the start of the test.
+                if self.simulation_dt == None:
+                    self.simulation_dt = date
+
+                # Done message has the risk report, so we yield before exiting.
+                if date == 'DONE':
+                    for event in snapshot:
+                        yield event.perf_message
 
                 # We're still in the warmup period.  Use the event to
-                # update our universe, but don't start a snapshot or
-                # pass anything to handle_data.  Discard any
-                # perf messages.
-                if event.dt != 'DONE' and event.dt < self.algo_start:
-                    self.update_universe(event)
-                    if event.perf_message:
-                        log.info("Discarding perf message because we're in warmup.")
-                    continue
+                # update our universe, but don't yield any perf messages,
+                # and don't send a snapshot to handle_data.
+                elif date < self.algo_start:
+                    for event in snapshot:
+                        del event['perf_message']
+                        self.update_universe(event)
 
-                # Yield any perf messages received to be relayed back to
-                # the browser.
-
-                if event.perf_message:
-                    yield event.perf_message
-                    del event['perf_message']
-
-                if event.dt == "DONE":
-                    if self.this_snapshot_dt:
-                        # StopIteration happened mid-snapshot, so we
-                        # have a universe snapshot that is not yet
-                        # processed by the algorithm.
-                        self.simulate_current_snapshot()
-
-                    # Break out of the loop, causing us to raise
-                    # StopIteration This needs to be outside the check
-                    # on self.this_snapshot_dt or else getting a DONE
-                    # immediately after a snapshot finishes will cause
-                    # type errors.
-                    break
-
-                # This should only happen for the first event we run.
-                if self.simulation_dt == None:
-                    self.simulation_dt = event.dt
-
-                # ======================
-                # Time Compression Logic
-                # ======================
-
-                if self.this_snapshot_dt != None:
-                    self.update_current_snapshot(event)
-
-                # The algorithm has been missing events because it took
-                # too long processing.  Update the universe with data from
-                # this event, then check if enough time has passed that we
-                # can start a new snapshot.
+                # The algo has taken so long to process events that
+                # its simulated time is later than the event time.
+                # Update the universe and yield any perf messages
+                # encountered, but don't call handle_data.
+                elif date < self.simulation_dt:
+                    for event in snapshot:
+                        # Only yield if we have something interesting to say.
+                        if event.perf_message != None:
+                            yield event.perf_message
+                        # Delete the message before updating so we don't send it
+                        # to the user.
+                        del event['perf_message']
+                        self.update_universe(event)
+                        
+                # Regular snapshot.  Update the universe and send a snapshot
+                # to handle data.
                 else:
-                    self.update_universe(event)
-                    if event.dt >= self.simulation_dt:
-                        self.this_snapshot_dt = event.dt
-
-
-
-    def update_current_snapshot(self, event):
-        """
-        Update our current snapshot of the universe. If event.dt doesn't
-        match our current snapshot's dt, we simulate the current snapshot
-        before processing the event.
-        """
-        # The new event matches our snapshot dt. Just update the
-        # universe and move on.
-        if event.dt == self.this_snapshot_dt:
-            self.update_universe(event)
-
-        # The new event does not match our snapshot.
-        else:
-            self.simulate_current_snapshot()
-
-            # Once we've finished simulating the old snapshot,
-            # we can update the universe with the new event.
-            self.update_universe(event)
-
-            # The current event is later than the simulation time,
-            # which means the algorithm finished quickly enough to
-            # receive the new event.  Start a new snapshot with this
-            # event's dt.
-            if event.dt >= self.simulation_dt:
-                self.this_snapshot_dt = event.dt
-
-            # The algorithm spent enough time processing that it
-            # missed the new event. Wait to start a new snapshot until
-            # the events catch up to the algo's simulated dt.
-            else:
-                self.this_snapshot_dt = None
-
-    def simulate_current_snapshot(self):
-        """
-        Run the user's algo against our current snapshot and update the algo's
-        simulated time.
-        """
-        start_tic = datetime.now()
-        self.algo.handle_data(self.universe)
-        stop_tic = datetime.now()
-
-        # How long did you take?
-        delta = stop_tic - start_tic
-
-        # Update the simulation time.
-        self.simulation_dt = self.this_snapshot_dt + delta
+                    for event in snapshot:
+                        # Only yield if we have something interesting to say.
+                        if event.perf_message != None:
+                            yield event.perf_message
+                        del event['perf_message']
+                        
+                        self.update_universe(event)
+                    
+                    # Send the current state of the universe to the user's algo.
+                    self.simulate_snapshot(date)
 
     def update_universe(self, event):
         """
@@ -341,3 +296,34 @@ class AlgorithmSimulator(object):
         for field in event.keys():
             self.universe[event.sid][field] = event[field]
 
+    # Ping every 10 seconds. Timeout after 9 pings.
+    @heartbeat(10, 9, self.handle_simulation_ping) 
+    def simulate_snapshot(self, date):
+        """
+        Run the user's algo against our current snapshot and update
+        the algo's simulated time.
+        """
+        start_tic = datetime.now()
+        
+        self.algo.handle_data(self.universe)
+        stop_tic = datetime.now()
+        
+        # How long did you take?
+        delta = stop_tic - start_tic
+
+        # Update the simulation time.
+        self.simulation_dt = date + delta
+
+    def handle_init_timeout(self, signum, frame):
+        """
+        Handler method for initialize timeout.
+        """
+        log.error("Algorithm timed out during initialize.")
+        raise AlgoTimeoutException("More than 5 seconds in initialize.")
+
+    def handle_simulation_ping(self, frame):
+        """
+        Frame handler for decorated simulate_snapshot method.
+        """
+
+    

From 58b027798a771178cad45be77f5afba8d4c5d84f Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sat, 18 Aug 2012 17:08:41 -0400
Subject: [PATCH 3/8] added timeout and heartbeat decorators

---
 zipline/utils/timeout.py | 102 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 102 insertions(+)
 create mode 100644 zipline/utils/timeout.py

diff --git a/zipline/utils/timeout.py b/zipline/utils/timeout.py
new file mode 100644
index 00000000..02b0903d
--- /dev/null
+++ b/zipline/utils/timeout.py
@@ -0,0 +1,102 @@
+import signal
+
+from pprint import pprint as pp
+from numbers import Number
+from logbook import Logger
+
+class Timeout(Exception):
+    
+    def __init__(self, frame):
+        self.frame = frame
+        
+
+class timeout(object):
+    """
+    Decorator to make a function raise TimeoutException if it spends
+    more than a specified number of seconds executing.
+    """
+
+    def __init__(self, seconds):
+        self.seconds = seconds
+        assert isinstance(seconds, Number), "Failed to specify a timeout."
+        assert seconds > 0, "Timeout must be greater than 0"
+
+    def handler(self, signum, frame):
+        raise Timeout(frame)
+
+    def __call__(self, fn):
+
+        def wrapped(*args, **kwargs):
+            # Set the alarm.
+            signal.signal(signal.SIGALRM, self.handler)
+            signal.alarm(self.seconds)
+            try:
+                outval = fn(*args, **kwargs)
+
+            # Deactivate the alarm once we're done so that the
+            # decorator doesn't have unexpected side-effects later.
+            # Note that this will still raise TimeoutException if the
+            # call to fn takes too long.
+            finally:
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, signal.SIG_DFL)
+
+            # Return the value of fn if it finished before the alarm.  This
+            # won't execute if the Timeout was raised.
+            return outval
+        return wrapped
+    
+class heartbeat(object):
+    """
+    Decorator to perform pseudo-heartbeat checks on a single-threaded
+    function. Calls frame_handler on the current stack frame of the 
+    decorated function every ``interval`` seconds.  After ``max_interval``
+    intervals, raises MaxHeartBeats
+    """
+
+    def __init__(self, interval, max_intervals, frame_handler=None):
+        self.count = 0
+        self.interval = interval
+        self.max_intervals = max_intervals
+        self.frame_handler = frame_handler
+        
+    def handler(self, signum, frame):
+        self.count += 1
+        if self.frame_handler:
+            self.frame_handler(frame)
+            
+        if self.count > self.max_intervals:
+            raise Timeout(frame)
+
+    def __call__(self, fn):
+        def wrapped(*args, **kwargs):
+            # Set a timer to call our handler every N seconds. 
+            signal.signal(signal.SIGALRM, self.handler)
+            signal.setitimer(signal.ITIMER_REAL, self.interval, self.interval)
+            try:
+                outval = fn(*args, **kwargs)
+
+            # Deactivate the timer once we're done so that the
+            # decorator doesn't have unexpected side-effects later.
+            finally:
+                signal.setitimer(signal.ITIMER_REAL, 0, 0)
+                signal.signal(signal.SIGALRM, signal.SIG_DFL)
+
+            # Return the value of fn if it finished without tripping
+            # an exception.  This won't execute if the Timeout or any
+            # other exception was raised by self.handle.
+            return outval
+        return wrapped
+   
+if __name__ == "__main__":
+    import time
+    
+    def pframe_g(frame):
+        print frame.f_globals
+
+    @heartbeat(1, 10, pframe_g)
+    def foo():
+        for i in xrange(10000):
+            time.sleep(.1)
+            print i
+    foo()

From 18a1c88145d0d1e472c02918c275c6325b40633b Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sat, 18 Aug 2012 21:10:28 -0400
Subject: [PATCH 4/8] fix import signal

---
 zipline/gens/tradesimulation.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 6a5e88b2..14669195 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -1,4 +1,4 @@
-
+import signal
 from logbook import Logger, Processor
 
 from datetime import datetime, timedelta
@@ -6,7 +6,7 @@ from numbers import Integral
 from itertools import groupby
 
 from zipline import ndict
-from zipline.utils import heartbeat
+from zipline.utils.timeout import heartbeat
 
 from zipline.gens.transform import StatefulTransform
 from zipline.finance.trading import TransactionSimulator
@@ -225,7 +225,7 @@ class AlgorithmSimulator(object):
         # snapshot time to any log record generated.
 
         with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''):
-            
+
             #Set an alarm to go off if initialize takes more than 5 seconds.
             signal.signal(signal.SIGALRM, self.handle_init_timeout)
             signal.alarm(5)
@@ -297,7 +297,7 @@ class AlgorithmSimulator(object):
             self.universe[event.sid][field] = event[field]
 
     # Ping every 10 seconds. Timeout after 9 pings.
-    @heartbeat(10, 9, self.handle_simulation_ping) 
+    # @heartbeat(10, 9, self.handle_simulation_ping) 
     def simulate_snapshot(self, date):
         """
         Run the user's algo against our current snapshot and update
@@ -325,5 +325,5 @@ class AlgorithmSimulator(object):
         """
         Frame handler for decorated simulate_snapshot method.
         """
-
+        print 'foo'
     

From 0ef89ee7049cbf37081ffbd254dcc18ca73d1047 Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sat, 18 Aug 2012 22:04:21 -0400
Subject: [PATCH 5/8] fix algo_dt logging

---
 zipline/gens/tradesimulation.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 14669195..8784a4fc 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -303,6 +303,10 @@ class AlgorithmSimulator(object):
         Run the user's algo against our current snapshot and update
         the algo's simulated time.
         """
+        # Needs to be set so that we inject the proper date into algo
+        # log/print lines.
+        self.snapshot_dt = date
+
         start_tic = datetime.now()
         
         self.algo.handle_data(self.universe)

From b5cb08ceedbfdeb4c91dde66cc729cfae3609314 Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sun, 19 Aug 2012 02:21:43 -0400
Subject: [PATCH 6/8] heartbeat and timeout usable as context managers

---
 zipline/utils/timeout.py | 97 +++++++++++++++++++++++++---------------
 1 file changed, 60 insertions(+), 37 deletions(-)

diff --git a/zipline/utils/timeout.py b/zipline/utils/timeout.py
index 02b0903d..ae844558 100644
--- a/zipline/utils/timeout.py
+++ b/zipline/utils/timeout.py
@@ -1,84 +1,109 @@
 import signal
 
+from functools import wraps
+
 from pprint import pprint as pp
 from numbers import Number
 from logbook import Logger
 
 class Timeout(Exception):
     
-    def __init__(self, frame):
+    def __init__(self, frame, message=''):
         self.frame = frame
+        self.message = message
         
-
 class timeout(object):
     """
-    Decorator to make a function raise TimeoutException if it spends
-    more than a specified number of seconds executing.
+    Utility to make a function raise TimeoutException if it spends
+    more than a specified number of seconds executing. Can be used
+    as a decorator to apply a static timeout to a function, or as 
+    a context manager to dynamically add a timeout to a code block.
     """
 
-    def __init__(self, seconds):
+    def __init__(self, seconds, message=''):
         self.seconds = seconds
+        self.message = message
         assert isinstance(seconds, Number), "Failed to specify a timeout."
         assert seconds > 0, "Timeout must be greater than 0"
 
     def handler(self, signum, frame):
-        raise Timeout(frame)
+        raise Timeout(frame, self.message)
 
     def __call__(self, fn):
-
-        def wrapped(*args, **kwargs):
+        
+        @wraps(fn)
+        def call_fn_with_timeout(*args, **kwargs):
             # Set the alarm.
             signal.signal(signal.SIGALRM, self.handler)
-            signal.alarm(self.seconds)
+            signal.setitimer(signal.ITIMER_REAL, self.seconds, 0)
             try:
                 outval = fn(*args, **kwargs)
 
             # Deactivate the alarm once we're done so that the
             # decorator doesn't have unexpected side-effects later.
-            # Note that this will still raise TimeoutException if the
+            # Note that this will still raise Timeout if the
             # call to fn takes too long.
             finally:
-                signal.alarm(0)
+                signal.setitimer(signal.ITIMER_REAL, 0, 0)
                 signal.signal(signal.SIGALRM, signal.SIG_DFL)
 
             # Return the value of fn if it finished before the alarm.  This
             # won't execute if the Timeout was raised.
             return outval
-        return wrapped
+        return call_fn_with_timeout
+
+    def __enter__(self):
+        # Set the alarm on entrance.
+        signal.signal(signal.SIGALRM, self.handler)
+        signal.setitimer(signal.ITIMER_REAL, self.seconds, 0)
+            
+    def __exit__(self, type, value, traceback):
+        # Deactivate the alarm on exit. This will re-raise
+        # any exceptions raised inside the with block.
+        signal.signal(signal.SIGALRM, self.handler)
+        signal.setitimer(signal.ITIMER_REAL, self.seconds, 0)
     
 class heartbeat(object):
     """
-    Decorator to perform pseudo-heartbeat checks on a single-threaded
+    Utility to perform pseudo-heartbeat checks on a single-threaded
     function. Calls frame_handler on the current stack frame of the 
-    decorated function every ``interval`` seconds.  After ``max_interval``
-    intervals, raises MaxHeartBeats
+    wrapped function every ``interval`` seconds.  After ``max_interval``
+    intervals, raises Timeout.  Can be used either as a decorator or
+    a context manager.
     """
+    def __init__(self, 
+                 interval, 
+                 max_intervals, 
+                 frame_handler=None, 
+                 timeout_message=''):
 
-    def __init__(self, interval, max_intervals, frame_handler=None):
-        self.count = 0
         self.interval = interval
         self.max_intervals = max_intervals
         self.frame_handler = frame_handler
+        self.timeout_message = timeout_message
+        self.count = 0
         
     def handler(self, signum, frame):
         self.count += 1
         if self.frame_handler:
-            self.frame_handler(frame)
+            self.frame_handler(self.count, frame)
             
-        if self.count > self.max_intervals:
-            raise Timeout(frame)
+        if self.count >= self.max_intervals:
+            raise Timeout(frame, self.timeout_message)
 
     def __call__(self, fn):
-        def wrapped(*args, **kwargs):
-            # Set a timer to call our handler every N seconds. 
+
+        @wraps(fn)
+        def call_fn_with_heartbeat(*args, **kwargs):
+            # Set a timer to call our handler every ``interval`` seconds. 
             signal.signal(signal.SIGALRM, self.handler)
             signal.setitimer(signal.ITIMER_REAL, self.interval, self.interval)
             try:
                 outval = fn(*args, **kwargs)
 
-            # Deactivate the timer once we're done so that the
-            # decorator doesn't have unexpected side-effects later.
             finally:
+                # Deactivate the timer once we're done so that the
+                # decorator doesn't have unexpected side-effects later.
                 signal.setitimer(signal.ITIMER_REAL, 0, 0)
                 signal.signal(signal.SIGALRM, signal.SIG_DFL)
 
@@ -86,17 +111,15 @@ class heartbeat(object):
             # an exception.  This won't execute if the Timeout or any
             # other exception was raised by self.handle.
             return outval
-        return wrapped
-   
-if __name__ == "__main__":
-    import time
+        return call_fn_with_heartbeat
     
-    def pframe_g(frame):
-        print frame.f_globals
-
-    @heartbeat(1, 10, pframe_g)
-    def foo():
-        for i in xrange(10000):
-            time.sleep(.1)
-            print i
-    foo()
+    def __enter__(self):
+        # Set a timer to call our handler every N seconds. 
+        signal.signal(signal.SIGALRM, self.handler)
+        signal.setitimer(signal.ITIMER_REAL, self.interval, self.interval)
+        
+    def __exit__(self, type, value, traceback):
+        # Turn off the timer on exit.  This will re-raise any exception raised
+        # during execution of the with-block
+        signal.setitimer(signal.ITIMER_REAL, 0, 0)
+        signal.signal(signal.SIGALRM, signal.SIG_DFL)

From 69ac68af2ee2d0010f2dbab22f7413752466da9c Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sun, 19 Aug 2012 02:22:16 -0400
Subject: [PATCH 7/8] heartbeating for handle_data

---
 tests/test_exception_handling.py | 35 +++++++++++++++++++-
 zipline/gens/tradesimulation.py  | 57 +++++++++++++++-----------------
 zipline/test_algorithms.py       | 50 ++++++++++++++++++++++++++++
 3 files changed, 110 insertions(+), 32 deletions(-)

diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py
index d1561837..ac6339bc 100644
--- a/tests/test_exception_handling.py
+++ b/tests/test_exception_handling.py
@@ -3,7 +3,8 @@ import zmq
 from unittest2 import TestCase
 from collections import defaultdict
 
-from zipline.test_algorithms import ExceptionAlgorithm, DivByZeroAlgorithm
+from zipline.test_algorithms import ExceptionAlgorithm, DivByZeroAlgorithm, \
+    InitializeTimeoutAlgorithm, TooMuchProcessingAlgorithm
 from zipline.finance.trading import SIMULATION_STYLE
 from zipline.core.devsimulator import AddressAllocator
 from zipline.lines import SimulatedTrading
@@ -143,3 +144,35 @@ class ExceptionTestCase(TestCase):
         # make sure our path shortening is working
         self.assertEqual(payload['stack'][0]['filename'], '/zipline/lines.py')
         self.assertEqual(payload['stack'][-1]['filename'], '/zipline/test_algorithms.py')
+
+    def test_initialize_timeout(self):
+        
+        self.zipline_test_config['algorithm'] = \
+            InitializeTimeoutAlgorithm(
+                self.zipline_test_config['sid']
+            )
+        
+        zipline = SimulatedTrading.create_test_zipline(
+            **self.zipline_test_config
+        )
+        output, _ = drain_zipline(self, zipline)
+        self.assertEqual(output[-1]['prefix'], 'EXCEPTION')
+        payload = output[-1]['payload']
+        self.assertEqual(payload['name'],'Timeout')
+        self.assertEqual(payload['message'], 'Call to initialize timed out')
+
+#     def test_heartbeat(self):
+        
+#         self.zipline_test_config['algorithm'] = \
+#             TooMuchProcessingAlgorithm(
+#                 self.zipline_test_config['sid']
+#             )
+#         zipline = SimulatedTrading.create_test_zipline(
+#             **self.zipline_test_config
+#         )
+#         output, _ = drain_zipline(self, zipline)
+#         self.assertEqual(output[-1]['prefix'], 'EXCEPTION')
+#         payload = output[-1]['payload']
+#         self.assertEqual(payload['name'],'Timeout')
+#         self.assertEqual(payload['message'], 'Too much time spent in handle_data call')
+        
diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 8784a4fc..811f99ac 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -6,7 +6,7 @@ from numbers import Integral
 from itertools import groupby
 
 from zipline import ndict
-from zipline.utils.timeout import heartbeat
+from zipline.utils.timeout import timeout, heartbeat, Timeout
 
 from zipline.gens.transform import StatefulTransform
 from zipline.finance.trading import TransactionSimulator
@@ -16,8 +16,10 @@ from zipline.gens.utils import hash_args
 
 log = Logger('Trade Simulation')
 
-class AlgoTimeoutException(Exception):
-    pass
+# TODO: make these arguments rather than global constants
+INIT_TIMEOUT = 5
+HEARTBEAT_INTERVAL = 1 # seconds
+MAX_HEARTBEAT_INTERVALS = 15
 
 class TradeSimulationClient(object):
     """
@@ -145,6 +147,21 @@ class AlgorithmSimulator(object):
         self.algolog = Logger("AlgoLog")
         self.algo.set_logger(self.algolog)
 
+        # Handler for heartbeats during calls to handle_data.
+        def log_heartbeats(beat_count, stackframe):
+            t = beat_count * HEARTBEAT_INTERVAL
+            warning = "handle_data has been processing for %i seconds" %t
+            self.algolog.warn(warning)
+
+        # Context manager that calls log_heartbeats every HEARTBEAT_INTERVAL
+        # seconds, raising an exception after MAX_HEARTBEATS
+        self.heartbeat_monitor = heartbeat(
+            HEARTBEAT_INTERVAL, 
+            MAX_HEARTBEAT_INTERVALS, 
+            frame_handler=log_heartbeats,
+            timeout_message="Too much time spent in handle_data call"
+        )
+
         # ==============
         # Snapshot Setup
         # ==============
@@ -223,17 +240,11 @@ class AlgorithmSimulator(object):
         # Capture any output of this generator to stdout and pipe it
         # to a logbook interface.  Also inject the current algo
         # snapshot time to any log record generated.
-
         with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''):
 
-            #Set an alarm to go off if initialize takes more than 5 seconds.
-            signal.signal(signal.SIGALRM, self.handle_init_timeout)
-            signal.alarm(5)
-            # Call the user's initialize method.
-            self.algo.initialize()
-            # Deactivate the alarm.
-            signal.alarm(0)
-            signal.signal(signal.SIGALRM, signal.SIG_DFL)
+            # Call user's initialize method with a timeout.
+            with timeout(INIT_TIMEOUT, message="Call to initialize timed out"):
+                self.algo.initialize()
 
             # Group together events with the same dt field. This depends on the
             # events already being sorted.
@@ -243,7 +254,7 @@ class AlgorithmSimulator(object):
                 # This should only occur once, at the start of the test.
                 if self.simulation_dt == None:
                     self.simulation_dt = date
-
+                    
                 # Done message has the risk report, so we yield before exiting.
                 if date == 'DONE':
                     for event in snapshot:
@@ -296,8 +307,6 @@ class AlgorithmSimulator(object):
         for field in event.keys():
             self.universe[event.sid][field] = event[field]
 
-    # Ping every 10 seconds. Timeout after 9 pings.
-    # @heartbeat(10, 9, self.handle_simulation_ping) 
     def simulate_snapshot(self, date):
         """
         Run the user's algo against our current snapshot and update
@@ -308,8 +317,8 @@ class AlgorithmSimulator(object):
         self.snapshot_dt = date
 
         start_tic = datetime.now()
-        
-        self.algo.handle_data(self.universe)
+        with self.heartbeat_monitor:
+            self.algo.handle_data(self.universe)
         stop_tic = datetime.now()
         
         # How long did you take?
@@ -317,17 +326,3 @@ class AlgorithmSimulator(object):
 
         # Update the simulation time.
         self.simulation_dt = date + delta
-
-    def handle_init_timeout(self, signum, frame):
-        """
-        Handler method for initialize timeout.
-        """
-        log.error("Algorithm timed out during initialize.")
-        raise AlgoTimeoutException("More than 5 seconds in initialize.")
-
-    def handle_simulation_ping(self, frame):
-        """
-        Frame handler for decorated simulate_snapshot method.
-        """
-        print 'foo'
-    
diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py
index a7881fa8..f6bdfd4d 100644
--- a/zipline/test_algorithms.py
+++ b/zipline/test_algorithms.py
@@ -221,6 +221,56 @@ class DivByZeroAlgorithm():
     def get_sid_filter(self):
         return [self.sid]
 
+class InitializeTimeoutAlgorithm():
+    def __init__(self, sid):
+        self.sid = sid
+        self.incr = 0
+
+    def initialize(self):
+        import time
+        from zipline.gens.tradesimulation import INIT_TIMEOUT
+        time.sleep(INIT_TIMEOUT + 1)
+
+    def set_order(self, order_callable):
+        pass
+
+    def set_logger(self, logger):
+        pass
+
+    def set_portfolio(self, portfolio):
+        pass
+
+    def handle_data(self, data):
+        pass
+
+    def get_sid_filter(self):
+        return [self.sid]
+    
+class TooMuchProcessingAlgorithm():
+    def __init__(self, sid):
+        self.sid = sid
+
+    def initialize(self):
+        pass
+
+    def set_order(self, order_callable):
+        pass
+
+    def set_logger(self, logger):
+        pass
+
+    def set_portfolio(self, portfolio):
+        pass
+
+    def handle_data(self, data):
+        # Unless we're running on some sort of
+        # supercomputer this will hit timeout.
+        for i in xrange(100000000):
+            self.foo = i
+
+    def get_sid_filter(self):
+        return [self.sid]
+    
 class TimeoutAlgorithm():
 
     def __init__(self, sid):

From 8b67d6a45c4fb4a17ed6b42713c4aca968e0b0f3 Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sun, 19 Aug 2012 15:00:09 -0400
Subject: [PATCH 8/8] fix __exit__ for timeout

---
 zipline/utils/timeout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zipline/utils/timeout.py b/zipline/utils/timeout.py
index ae844558..0310a208 100644
--- a/zipline/utils/timeout.py
+++ b/zipline/utils/timeout.py
@@ -61,7 +61,7 @@ class timeout(object):
         # Deactivate the alarm on exit. This will re-raise
         # any exceptions raised inside the with block.
         signal.signal(signal.SIGALRM, self.handler)
-        signal.setitimer(signal.ITIMER_REAL, self.seconds, 0)
+        signal.setitimer(signal.ITIMER_REAL, 0, 0)
     
 class heartbeat(object):
     """