diff --git a/tests/data/test_us_equity_pricing.py b/tests/data/test_us_equity_pricing.py index cbbc428c..e53cd7a9 100644 --- a/tests/data/test_us_equity_pricing.py +++ b/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): diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index c24005d0..2ce0cb2d 100644 --- a/tests/test_algorithm.py +++ b/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, diff --git a/tests/test_finance.py b/tests/test_finance.py index 7f307994..1de69f74 100644 --- a/tests/test_finance.py +++ b/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() ) diff --git a/zipline/data/bundles/core.py b/zipline/data/bundles/core.py index 9d1bf611..873d9ca3 100644 --- a/zipline/data/bundles/core.py +++ b/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 diff --git a/zipline/data/us_equity_loader.py b/zipline/data/us_equity_loader.py index 355f5954..c21c8557 100644 --- a/zipline/data/us_equity_loader.py +++ b/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( diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index 541c31f5..19e68e59 100644 --- a/zipline/data/us_equity_pricing.py +++ b/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 diff --git a/zipline/pipeline/loaders/equity_pricing_loader.py b/zipline/pipeline/loaders/equity_pricing_loader.py index aba34f20..813036c5 100644 --- a/zipline/pipeline/loaders/equity_pricing_loader.py +++ b/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( diff --git a/zipline/testing/core.py b/zipline/testing/core.py index 5c0f9e99..c661d7f0 100644 --- a/zipline/testing/core.py +++ b/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), ) diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index 76a60844..bd8bc415 100644 --- a/zipline/testing/fixtures.py +++ b/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())