diff --git a/tests/calendars/test_trading_calendar.py b/tests/calendars/test_trading_calendar.py index fa57225e..4745b513 100644 --- a/tests/calendars/test_trading_calendar.py +++ b/tests/calendars/test_trading_calendar.py @@ -56,13 +56,13 @@ class CalendarRegistrationTestCase(TestCase): dummy_cal = self.dummy_cal_type('DMY') # Try to register and retrieve the calendar - register_calendar(dummy_cal) + register_calendar('DMY', dummy_cal) retr_cal = get_calendar('DMY') self.assertEqual(dummy_cal, retr_cal) # Try to register again, expecting a name collision with self.assertRaises(CalendarNameCollision): - register_calendar(dummy_cal) + register_calendar('DMY', dummy_cal) # Deregister the calendar and ensure that it is removed deregister_calendar('DMY') @@ -76,7 +76,7 @@ class CalendarRegistrationTestCase(TestCase): real_nyse = get_calendar('NYSE') # Force a registration of the dummy NYSE - register_calendar(dummy_nyse, force=True) + register_calendar("NYSE", dummy_nyse, force=True) # Ensure that the dummy overwrote the real calendar retr_cal = get_calendar('NYSE') diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index bd8bc415..040ace99 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -37,7 +37,7 @@ from zipline.pipeline import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader from zipline.utils.calendars import ( get_calendar, -) + register_calendar) class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)): @@ -336,6 +336,8 @@ class WithAssetFinder(WithDefaultDateBounds): @classmethod def make_equity_info(cls): + register_calendar("TEST", get_calendar("NYSE"), force=True) + return make_simple_equity_info( cls.ASSET_FINDER_EQUITY_SIDS, cls.ASSET_FINDER_EQUITY_START_DATE, diff --git a/zipline/utils/calendars/calendar_utils.py b/zipline/utils/calendars/calendar_utils.py index 16f2ca4f..fc872601 100644 --- a/zipline/utils/calendars/calendar_utils.py +++ b/zipline/utils/calendars/calendar_utils.py @@ -28,24 +28,25 @@ def get_calendar(name): The desired calendar. """ if name not in _static_calendars: - if name == 'NYSE': + if name in ["NYSE", "NASDAQ", "BATS"]: cal = NYSEExchangeCalendar() - elif name == 'CME': + elif name in ["CME", "CBOT", "COMEX", "NYMEX"]: cal = CMEExchangeCalendar() + elif name in ["ICEUS", "NYFE"]: + cal = ICEExchangeCalendar() + elif name == "CFE": + cal = CFEExchangeCalendar() elif name == 'BMF': cal = BMFExchangeCalendar() elif name == 'LSE': cal = LSEExchangeCalendar() elif name == 'TSX': cal = TSXExchangeCalendar() - elif name == "ICE": - cal = ICEExchangeCalendar() - elif name == "CFE": - cal = CFEExchangeCalendar() + else: raise InvalidCalendarName(calendar_name=name) - register_calendar(cal) + register_calendar(name, cal) return _static_calendars[name] @@ -72,13 +73,15 @@ def clear_calendars(): _static_calendars.clear() -def register_calendar(calendar, force=False): +def register_calendar(name, calendar, force=False): """ Registers a calendar for retrieval by the get_calendar method. Parameters ---------- - calendar : TradingCalendar + name: str + The key with which to register this calendar. + calendar: TradingCalendar The calendar to be registered for retrieval. force : bool, optional If True, old calendars will be overwritten on a name collision. @@ -92,10 +95,10 @@ def register_calendar(calendar, force=False): # If we are forcing the registration, remove an existing calendar with the # same name. if force: - deregister_calendar(calendar.name) + deregister_calendar(name) # Check if we are already holding a calendar with the same name - if calendar.name in _static_calendars: - raise CalendarNameCollision(calendar_name=calendar.name) + if name in _static_calendars: + raise CalendarNameCollision(calendar_name=name) - _static_calendars[calendar.name] = calendar + _static_calendars[name] = calendar