Skip to content

Commit

Permalink
Merge pull request #1378 from quantopian/use-asset-type-to-key-fixtur…
Browse files Browse the repository at this point in the history
…e-calendars

TST: Key trading calendar fixture with Asset types
  • Loading branch information
ehebert committed Aug 8, 2016
2 parents fc11fd0 + f4891b0 commit a260fb1
Showing 1 changed file with 12 additions and 12 deletions.
24 changes: 12 additions & 12 deletions zipline/testing/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from ..utils.classproperty import classproperty
from ..utils.final import FinalMeta, final
from .core import tmp_asset_finder, make_simple_equity_info
from zipline.assets import Equity
from zipline.pipeline import SimplePipelineEngine
from zipline.pipeline.loaders.testing import make_seeded_random_loader
from zipline.utils.calendars import (
Expand Down Expand Up @@ -389,7 +390,7 @@ class WithTradingCalendars(object):
with that asset type.
"""
TRADING_CALENDAR_STRS = ('NYSE',)
TRADING_CALENDAR_FOR_ASSET_TYPE = {'equities': 'NYSE'}
TRADING_CALENDAR_FOR_ASSET_TYPE = {Equity: 'NYSE'}

# For backwards compatibility, exisitng tests and fixtures refer to
# `trading_calendar` with the assumption that the value is the NYSE
Expand All @@ -411,9 +412,6 @@ def init_class_fixtures(cls):
for asset_type, cal_str in iteritems(
cls.TRADING_CALENDAR_FOR_ASSET_TYPE):
calendar = get_calendar(cal_str)
setattr(cls,
'{0}_calendar'.format(asset_type),
calendar)
cls.trading_calendars[asset_type] = calendar


Expand Down Expand Up @@ -562,16 +560,17 @@ def init_class_fixtures(cls):

cls.trading_sessions = {}

for name, trading_calendar in iteritems(cls.trading_calendars):
for cal_str in cls.TRADING_CALENDAR_STRS:
trading_calendar = cls.trading_calendars[cal_str]
all_sessions = trading_calendar.all_sessions
start_loc = all_sessions.get_loc(cls.DATA_MIN_DAY, 'bfill')
end_loc = all_sessions.get_loc(cls.DATA_MAX_DAY, 'ffill')

sessions = all_sessions[start_loc:end_loc + 1]
# Set name for aliasing.
setattr(cls,
'{0}_sessions'.format(name.lower()), sessions)
cls.trading_sessions[name] = sessions
'{0}_sessions'.format(cal_str.lower()), sessions)
cls.trading_sessions[cal_str] = sessions


class WithTmpDir(object):
Expand Down Expand Up @@ -839,7 +838,7 @@ class which writes the minute bar data for use by a reader.

@classmethod
def make_equity_minute_bar_data(cls):
trading_calendar = cls.equities_calendar
trading_calendar = cls.trading_calendars[Equity]
return create_minute_bar_data(
trading_calendar.minutes_for_sessions_in_range(
cls.equity_minute_bar_days[0],
Expand All @@ -851,20 +850,21 @@ def make_equity_minute_bar_data(cls):
@classmethod
def init_class_fixtures(cls):
super(WithEquityMinuteBarData, cls).init_class_fixtures()
trading_calendar = cls.trading_calendars[Equity]
if cls.EQUITY_MINUTE_BAR_USE_FULL_CALENDAR:
days = cls.equites_calendar.all_execution_days
days = trading_calendar.all_execution_days
else:
first_session = cls.equities_calendar.minute_to_session_label(
first_session = trading_calendar.minute_to_session_label(
pd.Timestamp(cls.EQUITY_MINUTE_BAR_START_DATE)
)

if cls.EQUITY_MINUTE_BAR_LOOKBACK_DAYS > 0:
first_session = cls.equities_calendar.sessions_window(
first_session = trading_calendar.sessions_window(
first_session,
-1 * cls.EQUITY_MINUTE_BAR_LOOKBACK_DAYS
)[0]

days = cls.equities_calendar.sessions_in_range(
days = trading_calendar.sessions_in_range(
first_session,
cls.EQUITY_MINUTE_BAR_END_DATE
)
Expand Down

0 comments on commit a260fb1

Please sign in to comment.