MAINT: Use TradingCalendar objects for bundles (#1397)

* MAINT: Use TradingCalendar objects for bundles

Instead of trading days, opens, and closes, register now takes a
TradingCalendar object, along with a start_session and end_session. The
ingest function is now passed these values instead as well.

* Accept calendar name in addition to the actual object

* Updates bundles documentation for changes

* Fix typo in docs

* Use class formatting

* Force start_session and end_session within the bounds of the calendar

* Use UTC timestamps in test_core

* Document Trading Calendar API in appendix.rst
This commit is contained in:
Andrew Daniels
2016-08-17 13:37:07 -04:00
committed by GitHub
parent deac07d264
commit 440806ad60
8 changed files with 142 additions and 75 deletions
+25 -20
View File
@@ -30,7 +30,7 @@ from zipline.testing.predicates import (
)
from zipline.utils.cache import dataframe_cache
from zipline.utils.functional import apply
from zipline.utils.calendars import get_calendar
from zipline.utils.calendars import TradingCalendar, get_calendar
import zipline.utils.paths as pth
@@ -96,6 +96,8 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
daily_bar_writer,
adjustment_writer,
calendar,
start_session,
end_session,
cache,
show_progress,
output_dir):
@@ -111,20 +113,19 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
def test_ingest(self):
start = pd.Timestamp('2014-01-06', tz='utc')
end = pd.Timestamp('2014-01-10', tz='utc')
trading_days = get_calendar('NYSE').all_sessions
calendar = trading_days[trading_days.slice_indexer(start, end)]
minutes = get_calendar('NYSE').minutes_for_sessions_in_range(
calendar[0], calendar[-1]
)
calendar = get_calendar('NYSE')
sessions = calendar.sessions_in_range(start, end)
minutes = calendar.minutes_for_sessions_in_range(start, end)
sids = tuple(range(3))
equities = make_simple_equity_info(
sids,
calendar[0],
calendar[-1],
start,
end,
)
daily_bar_data = make_bar_data(equities, calendar)
daily_bar_data = make_bar_data(equities, sessions)
minute_bar_data = make_bar_data(equities, minutes)
first_split_ratio = 0.5
second_split_ratio = 0.1
@@ -141,13 +142,11 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
},
])
schedule = get_calendar('NYSE').schedule
@self.register(
'bundle',
calendar=calendar,
opens=schedule.market_open[calendar[0]:calendar[-1]],
closes=schedule.market_close[calendar[0]: calendar[-1]],
start_session=start,
end_session=end,
)
def bundle_ingest(environ,
asset_db_writer,
@@ -155,6 +154,8 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
daily_bar_writer,
adjustment_writer,
calendar,
start_session,
end_session,
cache,
show_progress,
output_dir):
@@ -165,7 +166,7 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
daily_bar_writer.write(daily_bar_data)
adjustment_writer.write(splits=splits)
assert_is_instance(calendar, pd.DatetimeIndex)
assert_is_instance(calendar, TradingCalendar)
assert_is_instance(cache, dataframe_cache)
assert_is_instance(show_progress, bool)
@@ -192,19 +193,19 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
columns,
calendar[0],
calendar[-1],
start,
end,
sids,
)
for actual_column, colname in zip(actual, columns):
assert_equal(
actual_column,
expected_bar_values_2d(calendar, equities, colname),
expected_bar_values_2d(sessions, equities, colname),
msg=colname,
)
adjustments_for_cols = bundle.adjustment_reader.load_adjustments(
columns,
calendar,
sessions,
pd.Index(sids),
)
for column, adjustments in zip(columns, adjustments_for_cols[:-1]):
@@ -263,7 +264,7 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
# register but do not ingest data
self.register('bundle', lambda *args: None)
ts = pd.Timestamp('2014')
ts = pd.Timestamp('2014', tz='UTC')
with assert_raises(ValueError) as e:
self.load('bundle', timestamp=ts, environ=self.environ)
@@ -291,13 +292,17 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
"""
if not self.bundles:
@self.register('bundle',
calendar=pd.DatetimeIndex([pd.Timestamp('2014')]))
calendar=get_calendar('NYSE'),
start_session=pd.Timestamp('2014', tz='UTC'),
end_session=pd.Timestamp('2014', tz='UTC'))
def _(environ,
asset_db_writer,
minute_bar_writer,
daily_bar_writer,
adjustment_writer,
calendar,
start_session,
end_session,
cache,
show_progress,
output_dir):
+11 -8
View File
@@ -27,9 +27,10 @@ class QuandlBundleTestCase(ZiplineTestCase):
symbols = 'AAPL', 'BRK_A', 'MSFT', 'ZEN'
asset_start = pd.Timestamp('2014-01', tz='utc')
asset_end = pd.Timestamp('2015-01', tz='utc')
calendar = bundles['quandl'].calendar
start_date = calendar[0]
end_date = calendar[-1]
bundle = bundles['quandl']
calendar = bundle.calendar
start_date = bundle.start_session
end_date = bundle.end_session
api_key = 'ayylmao'
columns = 'open', 'high', 'low', 'close', 'volume'
@@ -87,7 +88,9 @@ class QuandlBundleTestCase(ZiplineTestCase):
yield vs
# the first index our written data will appear in the files on disk
start_idx = self.calendar.get_loc(self.asset_start, 'ffill') + 1
start_idx = (
self.calendar.all_sessions.get_loc(self.asset_start, 'ffill') + 1
)
# convert an index into the raw dataframe into an index into the
# final data
@@ -215,11 +218,11 @@ class QuandlBundleTestCase(ZiplineTestCase):
assert_equal(equity.start_date, self.asset_start, msg=equity)
assert_equal(equity.end_date, self.asset_end, msg=equity)
cal = self.calendar
sessions = self.calendar.all_sessions
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
self.columns,
cal[cal.get_loc(self.asset_start, 'bfill')],
cal[cal.get_loc(self.asset_end, 'ffill')],
sessions[sessions.get_loc(self.asset_start, 'bfill')],
sessions[sessions.get_loc(self.asset_end, 'ffill')],
sids,
)
expected_pricing, expected_adjustments = self._expected_data(
@@ -229,7 +232,7 @@ class QuandlBundleTestCase(ZiplineTestCase):
adjustments_for_cols = bundle.adjustment_reader.load_adjustments(
self.columns,
cal,
sessions,
pd.Index(sids),
)
+9 -10
View File
@@ -20,11 +20,8 @@ class YahooBundleTestCase(WithResponses, ZiplineTestCase):
columns = 'open', 'high', 'low', 'close', 'volume'
asset_start = pd.Timestamp('2014-01-02', tz='utc')
asset_end = pd.Timestamp('2014-12-31', tz='utc')
trading_days = get_calendar('NYSE').all_sessions
calendar = trading_days[
(trading_days >= asset_start) &
(trading_days <= asset_end)
]
calendar = get_calendar('NYSE')
sessions = calendar.sessions_in_range(asset_start, asset_end)
@classmethod
def init_class_fixtures(cls):
@@ -157,11 +154,12 @@ class YahooBundleTestCase(WithResponses, ZiplineTestCase):
adjustments_callback,
)
cal = self.calendar
self.register(
'bundle',
yahoo_equities(self.symbols),
calendar=cal,
calendar=self.calendar,
start_session=self.asset_start,
end_session=self.asset_end,
)
zipline_root = self.enter_instance_context(tmp_dir()).path
@@ -181,10 +179,11 @@ class YahooBundleTestCase(WithResponses, ZiplineTestCase):
assert_equal(equity.start_date, self.asset_start, msg=equity)
assert_equal(equity.end_date, self.asset_end, msg=equity)
sessions = self.sessions
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
self.columns,
cal[cal.get_loc(self.asset_start, 'bfill')],
cal[cal.get_loc(self.asset_end, 'ffill')],
sessions[sessions.get_loc(self.asset_start, 'bfill')],
sessions[sessions.get_loc(self.asset_end, 'ffill')],
sids,
)
expected_pricing, expected_adjustments = self._expected_data()
@@ -192,7 +191,7 @@ class YahooBundleTestCase(WithResponses, ZiplineTestCase):
adjustments_for_cols = bundle.adjustment_reader.load_adjustments(
self.columns,
cal,
self.sessions,
pd.Index(sids),
)