From dfd9f2105d734c1e608bcbb8de789f6d6f490811 Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Thu, 31 Jan 2013 07:13:53 -0500
Subject: [PATCH] Changes test_source date_gen to only emit trading days.

Based on @fawce's work on record_vars branch.
---
 zipline/sources/test_source.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/zipline/sources/test_source.py b/zipline/sources/test_source.py
index f09174e2..f589392a 100644
--- a/zipline/sources/test_source.py
+++ b/zipline/sources/test_source.py
@@ -24,6 +24,7 @@ from datetime import datetime, timedelta
 import numpy as np
 
 from zipline.gens.utils import hash_args, create_trade
+from zipline.utils.tradingcalendar import trading_days
 
 
 def date_gen(start=datetime(2006, 6, 6, 12, tzinfo=pytz.utc),
@@ -33,12 +34,25 @@ def date_gen(start=datetime(2006, 6, 6, 12, tzinfo=pytz.utc),
     """
     Utility to generate a stream of dates.
     """
-    if repeats:
-        return (start + (i * delta)
-                for i in xrange(count)
-                for n in xrange(repeats))
-    else:
-        return (start + (i * delta) for i in xrange(count))
+    cur = start
+
+    # yield count trade events, all on trading days, and
+    # during trading hours.
+    # NB: Being inside of trading hours is currently dependent upon the
+    # count parameter being less than the number of trading minutes in a day
+    for i in xrange(count):
+        if repeats:
+            for j in xrange(repeats):
+                yield cur
+        else:
+            yield cur
+
+        cur = cur + delta
+        cur_midnight = cur.replace(hour=0, minute=0, second=0)
+        # skip over any non-trading days
+        if cur_midnight not in trading_days:
+            next_day_index = trading_days.searchsorted(cur_midnight)
+            cur_midnight = trading_days[next_day_index]
 
 
 def mock_prices(count):