From 800210fbb36a5cf24936c91292f731f9c673104d Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Thu, 17 Oct 2013 16:45:51 -0400
Subject: [PATCH] MAINT: Ensure that test sources only provide market days.

Instead of using all calendar days between start and end in test
sources, use the trading calendar for test sources.

Needed for an incoming refactoring of market open and close,
where the opens and closes are indexed by market days.
---
 zipline/finance/risk/cumulative.py |  8 +++---
 zipline/finance/trading.py         |  5 ++++
 zipline/utils/factory.py           | 40 +++++++++++++++++-------------
 3 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/zipline/finance/risk/cumulative.py b/zipline/finance/risk/cumulative.py
index ee0513c1..60bcd1e5 100644
--- a/zipline/finance/risk/cumulative.py
+++ b/zipline/finance/risk/cumulative.py
@@ -149,11 +149,9 @@ class RiskMetricsCumulative(object):
             hour=0, minute=0, second=0, microsecond=0
         )
 
-        all_trading_days = trading.environment.trading_days
-        mask = ((all_trading_days >= self.start_date) &
-                (all_trading_days <= self.end_date))
-
-        self.trading_days = all_trading_days[mask]
+        self.trading_days = trading.environment.days_in_range(
+            self.start_date,
+            self.end_date)
 
         last_day = normalize_date(sim_params.period_end)
         if last_day not in self.trading_days:
diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py
index 9b46aa50..cef7908d 100644
--- a/zipline/finance/trading.py
+++ b/zipline/finance/trading.py
@@ -162,6 +162,11 @@ class TradingEnvironment(object):
 
         return None
 
+    def days_in_range(self, start, end):
+        mask = ((self.trading_days >= start) &
+                (self.trading_days <= end))
+        return self.trading_days[mask]
+
     def next_open_and_close(self, start_date):
         """
         Given the start_date, returns the next open and close of
diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py
index 2c0b392e..db79ab7e 100644
--- a/zipline/utils/factory.py
+++ b/zipline/utils/factory.py
@@ -323,24 +323,23 @@ def create_test_df_source(sim_params=None, bars='daily'):
     if sim_params:
         index = sim_params.trading_days
     else:
+        if trading.environment is None:
+            trading.environment = trading.TradingEnvironment()
+
         start = pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc)
         end = pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)
-        index = pd.DatetimeIndex(
-            start=start,
-            end=end,
-            freq=freq
-        )
-        if bars == 'minute':
-            new_index = []
-            for i in index:
-                market_open = i.replace(hour=14,
-                                        minute=31)
-                market_close = i.replace(hour=21,
-                                         minute=0)
 
-                if i >= market_open and i <= market_close:
-                    new_index.append(i)
-            index = new_index
+        days = trading.environment.days_in_range(start, end)
+
+        if bars == 'daily':
+            index = days
+        if bars == 'minute':
+            index = pd.DatetimeIndex([], freq=freq)
+
+            for day in days:
+                day_index = trading.environment.market_minutes_for_day(day)
+                index = index.append(day_index)
+
     x = np.arange(1, len(index) + 1)
 
     df = pd.DataFrame(x, index=index, columns=[0])
@@ -355,7 +354,11 @@ def create_test_panel_source(sim_params=None):
     end = sim_params.last_close \
         if sim_params else pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)
 
-    index = pd.DatetimeIndex(start=start, end=end, freq=pd.datetools.day)
+    if trading.environment is None:
+        trading.environment = trading.TradingEnvironment()
+
+    index = trading.environment.days_in_range(start, end)
+
     price = np.arange(0, len(index))
     volume = np.ones(len(index)) * 1000
     arbitrary = np.ones(len(index))
@@ -376,7 +379,10 @@ def create_test_panel_ohlc_source(sim_params=None):
     end = sim_params.last_close \
         if sim_params else pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)
 
-    index = pd.DatetimeIndex(start=start, end=end, freq=pd.datetools.day)
+    if trading.environment is None:
+        trading.environment = trading.TradingEnvironment()
+
+    index = trading.environment.days_in_range(start, end)
     price = np.arange(0, len(index)) + 100
     high = price * 1.05
     low = price * 0.95