Merge pull request #1330 from quantopian/fix-daily-bar-writer

Fix daily bar writer

tests/data/test_us_equity_pricing.py

@@ -24,7 +24,6 @@ from numpy.testing import (
 )
 from pandas import (
     DataFrame,
-    DatetimeIndex,
     Timestamp,
 )
 from pandas.util.testing import assert_index_equal
@@ -46,6 +45,7 @@ from zipline.testing.fixtures import (
     WithBcolzEquityDailyBarReader,
     ZiplineTestCase,
 )
+from zipline.utils.calendars import get_calendar
 
 TEST_CALENDAR_START = Timestamp('2015-06-01', tz='UTC')
 TEST_CALENDAR_STOP = Timestamp('2015-06-30', tz='UTC')
@@ -180,9 +180,14 @@ class BcolzDailyBarTestCase(WithBcolzEquityDailyBarReader, ZiplineTestCase):
             result.attrs['calendar_offset'],
             expected_calendar_offset,
         )
+        cal = get_calendar(result.attrs['calendar_name'])
+        first_session = Timestamp(result.attrs['start_session_ns'], tz='UTC')
+        end_session = Timestamp(result.attrs['end_session_ns'], tz='UTC')
+        sessions = cal.sessions_in_range(first_session, end_session)
+
         assert_index_equal(
             self.sessions,
-            DatetimeIndex(result.attrs['calendar'], tz='UTC'),
+            sessions
         )
 
     def test_read_first_trading_day(self):

tests/test_algorithm.py

@@ -3424,9 +3424,10 @@ class TestEquityAutoClose(WithTmpDir, WithTradingCalendar, ZiplineTestCase):
                 frequency=frequency
             )
             path = self.tmpdir.getpath("testdaily.bcolz")
-            BcolzDailyBarWriter(path, dates, self.trading_calendar).write(
-                iteritems(trade_data_by_sid),
+            writer = BcolzDailyBarWriter(
+                path, self.trading_calendar, dates[0], dates[-1]
             )
+            writer.write(iteritems(trade_data_by_sid))
             reader = BcolzDailyBarReader(path)
             data_portal = DataPortal(
                 env.asset_finder, self.trading_calendar,

tests/test_finance.py

@@ -255,7 +255,8 @@ class FinanceTestCase(WithLogger,
                 }
 
                 path = os.path.join(tempdir.path, "testdata.bcolz")
-                BcolzDailyBarWriter(path, days, self.trading_calendar).write(
+                BcolzDailyBarWriter(path, self.trading_calendar, days[0],
+                                    days[-1]).write(
                     assets.items()
                 )
 

zipline/data/bundles/core.py

@@ -333,7 +333,9 @@ def _make_bundle_core():
                 )).path
                 daily_bar_writer = BcolzDailyBarWriter(
                     daily_bars_path,
-                    bundle.calendar,
+                    nyse_cal,
+                    bundle.calendar[0],
+                    bundle.calendar[-1]
                 )
                 # Do an empty write to ensure that the daily ctables exist
                 # when we create the SQLiteAdjustmentWriter below. The

zipline/data/us_equity_loader.py

@@ -387,7 +387,7 @@ class USEquityDailyHistoryLoader(USEquityHistoryLoader):
 
     @property
     def _calendar(self):
-        return self._reader._calendar
+        return self._reader._sessions
 
     def _array(self, dts, assets, field):
         return self._reader.load_raw_arrays(

zipline/data/us_equity_pricing.py

@@ -40,12 +40,11 @@ from numpy import (
 )
 from pandas import (
     DataFrame,
-    DatetimeIndex,
     read_csv,
     Timestamp,
     NaT,
     isnull,
-)
+    DatetimeIndex)
 from pandas.tslib import iNaT
 from six import (
     iteritems,
@@ -53,6 +52,7 @@ from six import (
     viewkeys,
 )
 
+from zipline.utils.calendars import get_calendar
 from zipline.utils.functional import apply
 from zipline.utils.preprocess import call
 from zipline.utils.input_validation import (
@@ -182,15 +182,19 @@ def to_ctable(raw_data, invalid_data_behavior):
 
 class BcolzDailyBarWriter(object):
     """
-    Class capable of writing daily OHLCV data to disk in a format that can be
-    read efficiently by BcolzDailyOHLCVReader.
+    Class capable of writing daily OHLCV data to disk in a format that can
+    be read efficiently by BcolzDailyOHLCVReader.
 
     Parameters
     ----------
     filename : str
         The location at which we should write our output.
-    sessions : pandas.DatetimeIndex
+    calendar : zipline.utils.calendar.trading_calendar
         Calendar to use to compute asset calendar offsets.
+    start_session: pd.Timestamp
+        Midnight UTC session label.
+    end_session: pd.Timestamp
+        Midnight UTC session label.
 
     See Also
     --------
@@ -204,9 +208,15 @@ class BcolzDailyBarWriter(object):
         'volume': float64,
     }
 
-    def __init__(self, filename, sessions, calendar):
+    def __init__(self, filename, calendar, start_session, end_session):
         self._filename = filename
-        self._sessions = sessions
+
+        assert calendar.is_session(start_session), "Start session is invalid!"
+        assert calendar.is_session(end_session), "End session is invalid!"
+
+        self._start_session = start_session
+        self._end_session = end_session
+
         self._calendar = calendar
 
     @property
@@ -300,7 +310,9 @@ class BcolzDailyBarWriter(object):
         }
 
         earliest_date = None
-        sessions = self._sessions
+        sessions = self._calendar.sessions_in_range(
+            self._start_session, self._end_session
+        )
 
         if assets is not None:
             @apply
@@ -363,10 +375,13 @@ class BcolzDailyBarWriter(object):
         full_table.attrs['first_trading_day'] = (
             earliest_date if earliest_date is not None else iNaT
         )
+
         full_table.attrs['first_row'] = first_row
         full_table.attrs['last_row'] = last_row
         full_table.attrs['calendar_offset'] = calendar_offset
-        full_table.attrs['calendar'] = sessions.asi8.tolist()
+        full_table.attrs['calendar_name'] = self._calendar.name
+        full_table.attrs['start_session_ns'] = self._start_session.value
+        full_table.attrs['end_session_ns'] = self._end_session.value
         full_table.flush()
         return full_table
 
@@ -387,6 +402,14 @@ class DailyBarReader(with_metaclass(ABCMeta)):
     def last_available_dt(self):
         pass
 
+    @abstractproperty
+    def trading_calendar(self):
+        """
+        Returns the zipline.utils.calendar.trading_calendar used to read
+         the data.  Can be None (if the writer didn't specify it).
+        """
+        pass
+
 
 class BcolzDailyBarReader(DailyBarReader):
     """
@@ -415,8 +438,12 @@ class BcolzDailyBarReader(DailyBarReader):
         Map from asset_id -> index of last row in the dataset with that id.
     calendar_offset : dict
         Map from asset_id -> calendar index of first row.
-    calendar : list[int64]
-        Calendar used to compute offsets, in asi8 format (ns since EPOCH).
+    start_session_ns: int
+        Epoch ns of the first session used in this dataset.
+    end_session_ns: int
+        Epoch ns of the last session used in this dataset.
+    calendar_name: str
+        String identifier of trading calendar used (ie, "NYSE").
 
     We use first_row and last_row together to quickly find ranges of rows to
     load when reading an asset's data into memory.
@@ -473,8 +500,21 @@ class BcolzDailyBarReader(DailyBarReader):
         return ctable(rootdir=maybe_table_rootdir, mode='r')
 
     @lazyval
-    def _calendar(self):
-        return DatetimeIndex(self._table.attrs['calendar'], tz='UTC')
+    def _sessions(self):
+        if 'calendar' in self._table.attrs.attrs:
+            # backwards compatibility with old formats, will remove
+            return DatetimeIndex(self._table.attrs['calendar'], tz='UTC')
+        else:
+            cal = get_calendar(self._table.attrs['calendar_name'])
+            start_session_ns = self._table.attrs['start_session_ns']
+            start_session = Timestamp(start_session_ns, tz='UTC')
+
+            end_session_ns = self._table.attrs['end_session_ns']
+            end_session = Timestamp(end_session_ns, tz='UTC')
+
+            sessions = cal.sessions_in_range(start_session, end_session)
+
+            return sessions
 
     @lazyval
     def _first_rows(self):
@@ -514,9 +554,16 @@ class BcolzDailyBarReader(DailyBarReader):
         except KeyError:
             return None
 
+    @lazyval
+    def trading_calendar(self):
+        if 'calendar_name' in self._table.attrs.attrs:
+            return get_calendar(self._table.attrs['calendar_name'])
+        else:
+            return None
+
     @property
     def last_available_dt(self):
-        return self._calendar[-1]
+        return self._sessions[-1]
 
     def _compute_slices(self, start_idx, end_idx, assets):
         """
@@ -562,8 +609,8 @@ class BcolzDailyBarReader(DailyBarReader):
 
     def load_raw_arrays(self, columns, start_date, end_date, assets):
         # Assumes that the given dates are actually in calendar.
-        start_idx = self._calendar.get_loc(start_date)
-        end_idx = self._calendar.get_loc(end_date)
+        start_idx = self._sessions.get_loc(start_date)
+        end_idx = self._sessions.get_loc(end_date)
         first_rows, last_rows, offsets = self._compute_slices(
             start_idx,
             end_idx,
@@ -607,8 +654,8 @@ class BcolzDailyBarReader(DailyBarReader):
 
         if day >= asset.end_date:
             # go back to one day before the asset ended
-            search_day = self._calendar[
-                self._calendar.searchsorted(asset.end_date) - 1
+            search_day = self._sessions[
+                self._sessions.searchsorted(asset.end_date) - 1
             ]
         else:
             search_day = day
@@ -620,9 +667,9 @@ class BcolzDailyBarReader(DailyBarReader):
                 return None
             if volumes[ix] != 0:
                 return search_day
-            prev_day_ix = self._calendar.get_loc(search_day) - 1
+            prev_day_ix = self._sessions.get_loc(search_day) - 1
             if prev_day_ix > -1:
-                search_day = self._calendar[prev_day_ix]
+                search_day = self._sessions[prev_day_ix]
             else:
                 return None
 
@@ -643,10 +690,10 @@ class BcolzDailyBarReader(DailyBarReader):
             or after the date range of the equity.
         """
         try:
-            day_loc = self._calendar.get_loc(day)
+            day_loc = self._sessions.get_loc(day)
         except:
             raise NoDataOnDate("day={0} is outside of calendar={1}".format(
-                day, self._calendar))
+                day, self._sessions))
         offset = day_loc - self._calendar_offsets[sid]
         if offset < 0:
             raise NoDataOnDate(
@@ -729,6 +776,10 @@ class PanelDailyBarReader(DailyBarReader):
     def last_available_dt(self):
         return self._calendar[-1]
 
+    @property
+    def trading_calendar(self):
+        return None
+
     def load_raw_arrays(self, columns, start_date, end_date, assets):
         columns = list(columns)
         cal = self._calendar

zipline/pipeline/loaders/equity_pricing_loader.py

@@ -40,7 +40,10 @@ class USEquityPricingLoader(PipelineLoader):
         self.raw_price_loader = raw_price_loader
         self.adjustments_loader = adjustments_loader
 
-        self._calendar = get_calendar("NYSE").all_sessions
+        cal = self.raw_price_loader.trading_calendar or \
+            get_calendar("NYSE")
+
+        self._all_sessions = cal.all_sessions
 
     @classmethod
     def from_files(cls, pricing_path, adjustments_path):
@@ -67,7 +70,7 @@ class USEquityPricingLoader(PipelineLoader):
         # known on day N is the data from day (N - 1), so we shift all query
         # dates back by a day.
         start_date, end_date = _shift_dates(
-            self._calendar, dates[0], dates[-1], shift=1,
+            self._all_sessions, dates[0], dates[-1], shift=1,
         )
         colnames = [c.name for c in columns]
         raw_arrays = self.raw_price_loader.load_raw_arrays(

zipline/testing/core.py

@@ -462,7 +462,9 @@ def create_daily_bar_data(sessions, sids):
 
 def write_daily_data(tempdir, sim_params, sids, trading_calendar):
     path = os.path.join(tempdir.path, "testdaily.bcolz")
-    BcolzDailyBarWriter(path, sim_params.sessions, trading_calendar).write(
+    BcolzDailyBarWriter(path, trading_calendar,
+                        sim_params.start_session,
+                        sim_params.end_session).write(
         create_daily_bar_data(sim_params.sessions, sids),
     )
 
@@ -612,7 +614,12 @@ def create_data_portal_from_trade_history(asset_finder, trading_calendar,
                                           tempdir, sim_params, trades_by_sid):
     if sim_params.data_frequency == "daily":
         path = os.path.join(tempdir.path, "testdaily.bcolz")
-        BcolzDailyBarWriter(path, sim_params.sessions, trading_calendar).write(
+        writer = BcolzDailyBarWriter(
+            path, trading_calendar,
+            sim_params.start_session,
+            sim_params.end_session
+        )
+        writer.write(
             trades_by_sid_to_dfs(trades_by_sid, sim_params.sessions),
         )
 

zipline/testing/fixtures.py

@@ -746,7 +746,7 @@ class WithBcolzEquityDailyBarReader(WithEquityDailyBarData, WithTmpDir):
         days = cls.equity_daily_bar_days
 
         cls.bcolz_daily_bar_ctable = t = getattr(
-            BcolzDailyBarWriter(p, days, cls.trading_calendar),
+            BcolzDailyBarWriter(p, cls.trading_calendar, days[0], days[-1]),
             cls._write_method_name,
         )(cls.make_equity_daily_bar_data())