mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-05 12:50:21 +08:00
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:
@@ -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):
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user