From 9bf4855b8cdf1d227c5dae189da8f6eb9060e852 Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Tue, 30 Jun 2015 11:30:13 -0400
Subject: [PATCH] MAINT: Move ZiplineAPI context so that it always wraps main
 loop.

Move the responsibility of wrapping the main simulation loop in the
ZiplineAPI context from the algorithm modules generator setup to the
main trade simulation loop, so that different methods of invoking loop
do not need to duplicate how the context is set.

To make it easier for internal implementation of handle_data to
transition off of calling the ZiplineAPI every bar, to only invoking the
context once per simulation.
---
 etc/requirements.txt            |  2 ++
 zipline/algorithm.py            | 15 +++++++--------
 zipline/gens/tradesimulation.py | 10 +++++++++-
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/etc/requirements.txt b/etc/requirements.txt
index 5c189006..81cd4320 100644
--- a/etc/requirements.txt
+++ b/etc/requirements.txt
@@ -31,3 +31,5 @@ cyordereddict==0.2.2
 
 # faster array ops.
 bottleneck==1.0.0
+
+contextlib2==0.4.0
diff --git a/zipline/algorithm.py b/zipline/algorithm.py
index 5307cbee..0a8b5e2e 100644
--- a/zipline/algorithm.py
+++ b/zipline/algorithm.py
@@ -510,15 +510,14 @@ class TradingAlgorithm(object):
                 self.sim_params.data_frequency,
             )
 
-        with ZiplineAPI(self):
-            # loop through simulated_trading, each iteration returns a
-            # perf dictionary
-            perfs = []
-            for perf in self.gen:
-                perfs.append(perf)
+        # loop through simulated_trading, each iteration returns a
+        # perf dictionary
+        perfs = []
+        for perf in self.gen:
+            perfs.append(perf)
 
-            # convert perf dict to pandas dataframe
-            daily_stats = self._create_daily_stats(perfs)
+        # convert perf dict to pandas dataframe
+        daily_stats = self._create_daily_stats(perfs)
 
         self.analyze(daily_stats)
 
diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 9f926ac4..54bb575c 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -13,9 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from contextlib2 import ExitStack
+
 from logbook import Logger, Processor
 from pandas.tslib import normalize_date
 
+from zipline.utils.api_support import ZiplineAPI
+
 from zipline.finance import trading
 from zipline.protocol import (
     BarData,
@@ -81,7 +85,11 @@ class AlgorithmSimulator(object):
 
         # inject the current algo
         # snapshot time to any log record generated.
-        with self.processor.threadbound():
+
+        with ExitStack() as stack:
+            stack.enter_context(self.processor.threadbound())
+            stack.enter_context(ZiplineAPI(self.algo))
+
             data_frequency = self.sim_params.data_frequency
 
             self._call_before_trading_start(mkt_open)