MAINT: Clean up set_algo_instance usage in TradingAlgorithm.

TradingAlgorithm always uses set_algo_instance in pairs of
set_algo_instance(self) and set_algo_instance(None).  Refactoring this to use a
context manager.

zipline/algorithm.py

@@ -34,7 +34,7 @@ from zipline.errors import (
 from zipline.finance.performance import PerformanceTracker
 from zipline.sources import DataFrameSource, DataPanelSource
 from zipline.utils.factory import create_simulation_parameters
-from zipline.utils.api_support import set_algo_instance, api_method
+from zipline.utils.api_support import ZiplineAPI, api_method
 from zipline.transforms.utils import StatefulTransform
 from zipline.finance.slippage import (
     VolumeShareSlippage,
@@ -195,12 +195,12 @@ class TradingAlgorithm(object):
         self.initialize(*args, **kwargs)
 
     def initialize(self, *args, **kwargs):
-        # store algo reference in global space
-        set_algo_instance(self)
-        try:
+        """
+        Call self._initialize with `self` made available to Zipline API
+        functions.
+        """
+        with ZiplineAPI(self):
             self._initialize(self)
-        finally:
-            set_algo_instance(None)
 
     def handle_data(self, data):
         if self.history_container:
@@ -402,10 +402,7 @@ class TradingAlgorithm(object):
         # create transforms and zipline
         self.gen = self._create_generator(sim_params)
 
-        # store algo reference in global space
-        set_algo_instance(self)
-
-        try:
+        with ZiplineAPI(self):
             # loop through simulated_trading, each iteration returns a
             # perf dictionary
             perfs = []
@@ -414,9 +411,6 @@ class TradingAlgorithm(object):
 
             # convert perf dict to pandas dataframe
             daily_stats = self._create_daily_stats(perfs)
-        finally:
-            # remove algo from global space
-            set_algo_instance(None)
 
         return daily_stats
 

zipline/utils/api_support.py

@@ -28,6 +28,29 @@ def set_algo_instance(algo):
     context.algorithm = algo
 
 
+class ZiplineAPI(object):
+    """
+    Context manager for making an algorithm instance available to zipline API
+    functions within a scoped block.
+    """
+
+    def __init__(self, algo_instance):
+        self.algo_instance = algo_instance
+
+    def __enter__(self):
+        """
+        Set the given algo instance, storing any previously-existing instance.
+        """
+        self.old_algo_instance = get_algo_instance()
+        set_algo_instance(self.algo_instance)
+
+    def __exit__(self, _type, _value, _tb):
+        """
+        Restore the algo instance stored in __enter__.
+        """
+        set_algo_instance(self.old_algo_instance)
+
+
 def api_method(f):
     # Decorator that adds the decorated class method as a callable
     # function (wrapped) to zipline.api