Merge pull request #1432 from quantopian/reindex-reader

ENH: Add a reader base which reindexes results.

tests/data/test_resample.py

@@ -16,7 +16,7 @@ from numbers import Real
 
 from nose_parameterized import parameterized
 from numpy.testing import assert_almost_equal
-from numpy import nan, array
+from numpy import nan, array, full
 import pandas as pd
 from pandas import DataFrame
 from six import iteritems
@@ -25,11 +25,14 @@ from zipline.data.resample import (
     minute_to_session,
     DailyHistoryAggregator,
     MinuteResampleSessionBarReader,
+    ReindexMinuteBarReader,
+    ReindexSessionBarReader,
 )
 
 from zipline.testing.fixtures import (
     WithEquityMinuteBarData,
     WithBcolzEquityMinuteBarReader,
+    WithBcolzEquityDailyBarReader,
     WithBcolzFutureMinuteBarReader,
     ZiplineTestCase,
 )
@@ -527,3 +530,124 @@ class TestResampleSessionBars(WithBcolzFutureMinuteBarReader,
                     assert_almost_equal(values[col], result,
                                         err_msg="sid={0} col={1} dt={2}".
                                         format(sid, col, dt))
+
+
+class TestReindexMinuteBars(WithBcolzEquityMinuteBarReader,
+                            ZiplineTestCase):
+
+    TRADING_CALENDAR_STRS = ('CME', 'NYSE')
+    TRADING_CALENDAR_PRIMARY_CAL = 'CME'
+
+    ASSET_FINDER_EQUITY_SIDS = 1, 2, 3
+
+    START_DATE = pd.Timestamp('2015-12-01', tz='UTC')
+    END_DATE = pd.Timestamp('2015-12-31', tz='UTC')
+
+    def test_load_raw_arrays(self):
+        reindex_reader = ReindexMinuteBarReader(
+            self.trading_calendar,
+            self.bcolz_equity_minute_bar_reader,
+            self.START_DATE,
+            self.END_DATE,
+        )
+        m_open, m_close = self.trading_calendar.open_and_close_for_session(
+            self.START_DATE)
+        outer_minutes = self.trading_calendar.minutes_in_range(m_open, m_close)
+        result = reindex_reader.load_raw_arrays(
+            OHLCV, m_open, m_close, [1, 2])
+
+        opens = DataFrame(data=result[0], index=outer_minutes,
+                          columns=[1, 2])
+        opens_with_price = opens.dropna()
+
+        self.assertEqual(
+            1440,
+            len(opens),
+            "The result should have 1440 bars, the number of minutes in a "
+            "trading session on the target calendar."
+        )
+
+        self.assertEqual(
+            390,
+            len(opens_with_price),
+            "The result, after dropping nans, should have 390 bars, the "
+            " number of bars in a trading session in the reader's calendar."
+        )
+
+        slicer = outer_minutes.slice_indexer(
+            end=pd.Timestamp('2015-12-01 14:30', tz='UTC'))
+
+        assert_almost_equal(
+            opens[1][slicer],
+            full(slicer.stop, nan),
+            err_msg="All values before the NYSE market open should be nan.")
+
+        slicer = outer_minutes.slice_indexer(
+            start=pd.Timestamp('2015-12-01 21:01', tz='UTC'))
+
+        assert_almost_equal(
+            opens[1][slicer],
+            full(slicer.stop - slicer.start, nan),
+            err_msg="All values after the NYSE market close should be nan.")
+
+        first_minute_loc = outer_minutes.get_loc(pd.Timestamp(
+            '2015-12-01 14:31', tz='UTC'))
+
+        # Spot check a value.
+        # The value is the autogenerated value from test fixtures.
+        assert_almost_equal(
+            10.0,
+            opens[1][first_minute_loc],
+            err_msg="The value for Equity 1, should be 10.0, at NYSE open.")
+
+
+class TestReindexSessionBars(WithBcolzEquityDailyBarReader,
+                             ZiplineTestCase):
+
+    TRADING_CALENDAR_STRS = ('CME', 'NYSE')
+    TRADING_CALENDAR_PRIMARY_CAL = 'CME'
+
+    ASSET_FINDER_EQUITY_SIDS = 1, 2, 3
+
+    # Dates are chosen to span Thanksgiving, which is not a Holiday on CME.
+    START_DATE = pd.Timestamp('2015-11-01', tz='UTC')
+    END_DATE = pd.Timestamp('2015-11-30', tz='UTC')
+
+    def test_load_raw_arrays(self):
+        reindex_reader = ReindexSessionBarReader(
+            self.trading_calendar,
+            self.bcolz_equity_daily_bar_reader,
+            self.START_DATE,
+            self.END_DATE,
+        )
+
+        outer_sessions = self.trading_calendar.sessions_in_range(
+            self.START_DATE, self.END_DATE)
+
+        result = reindex_reader.load_raw_arrays(
+            OHLCV, self.START_DATE, self.END_DATE, [1, 2])
+
+        opens = DataFrame(data=result[0], index=outer_sessions,
+                          columns=[1, 2])
+        opens_with_price = opens.dropna()
+
+        self.assertEqual(
+            21,
+            len(opens),
+            "The reindexed result should have 21 days, which is the number of "
+            "business days in 2015-11")
+        self.assertEqual(
+            20,
+            len(opens_with_price),
+            "The reindexed result after dropping nans should have 20 days, "
+            "because Thanksgiving is a NYSE holiday.")
+
+        # Thanksgiving, 2015-11-26.
+        # Is a holiday in NYSE, but not in CME.
+        tday_loc = outer_sessions.get_loc(pd.Timestamp('2015-11-26', tz='UTC'))
+
+        assert_almost_equal(
+            nan,
+            opens[1][tday_loc],
+            err_msg="2015-11-26 should be `nan`, since Thanksgiving is a "
+            "holiday in the reader's calendar.")