From 1752f78447ec98df6af9fd22eda8cb6148a2eaf9 Mon Sep 17 00:00:00 2001
From: fawce <fawce@quantopian.com>
Date: Thu, 2 May 2013 16:53:51 -0400
Subject: [PATCH] ENH: Allow algorithm to run past end of trading.environment
 history.

Work towards running an algorithm against 'live' data, which can't
be bound to the available benchmarks and treasuries, since the
benchmarks and treasury curves for that day won't be published
until that night.
---
 zipline/finance/trading.py      |  5 ++++-
 zipline/gens/tradesimulation.py | 13 +++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py
index b610ea65..c9dde661 100644
--- a/zipline/finance/trading.py
+++ b/zipline/finance/trading.py
@@ -75,7 +75,8 @@ class TradingEnvironment(object):
         self,
         load=None,
         bm_symbol='^GSPC',
-        exchange_tz="US/Eastern"
+        exchange_tz="US/Eastern",
+        max_date=None
     ):
         self.prev_environment = self
         self.trading_day_map = OrderedDict()
@@ -94,6 +95,8 @@ class TradingEnvironment(object):
         self.exchange_tz = exchange_tz
 
         for bm in self.benchmark_returns:
+            if max_date and bm.date > max_date:
+                break
             self.trading_day_map[bm.date] = bm
 
         self.first_trading_day = next(self.trading_day_map.iterkeys())
diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index dd409f05..09421d10 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -168,13 +168,8 @@ class AlgorithmSimulator(object):
                             yield daily_rollup
                             tp = self.algo.perf_tracker.todays_performance
                             tp.rollover()
-
                             if mkt_close < self.algo.perf_tracker.last_close:
-                                env = trading.environment
-                                _, mkt_close = \
-                                    env.next_open_and_close(
-                                        mkt_close
-                                    )
+                                mkt_close = self.get_next_close(mkt_close)
 
             risk_message = self.algo.perf_tracker.handle_simulation_end()
             yield risk_message
@@ -193,6 +188,12 @@ class AlgorithmSimulator(object):
             perf_message['intraday_perf']['recorded_vars'] = rvars
             return perf_message
 
+    def get_next_close(self, mkt_close):
+        if mkt_close >= trading.environment.last_trading_day:
+            return self.sim_params.last_close
+        else:
+            return trading.environment.next_open_and_close(mkt_close)[1]
+
     def update_universe(self, event):
         """
         Update the universe with new event information.