From 7803ec6e460f7ecbd91865d98240b4c4133c7ed9 Mon Sep 17 00:00:00 2001 From: Jean Bredeche Date: Wed, 10 Aug 2016 09:49:53 -0400 Subject: [PATCH] ENH: Add public API to register calendars by type --- tests/calendars/test_trading_calendar.py | 66 +++++++++++++++++------ zipline/utils/calendars/__init__.py | 3 +- zipline/utils/calendars/calendar_utils.py | 34 ++++++++++-- 3 files changed, 84 insertions(+), 19 deletions(-) diff --git a/tests/calendars/test_trading_calendar.py b/tests/calendars/test_trading_calendar.py index 4745b513..2fdbd2b2 100644 --- a/tests/calendars/test_trading_calendar.py +++ b/tests/calendars/test_trading_calendar.py @@ -19,7 +19,6 @@ from os.path import ( join, ) from unittest import TestCase -from collections import namedtuple import numpy as np import pandas as pd @@ -38,22 +37,40 @@ from zipline.utils.calendars import( register_calendar, deregister_calendar, get_calendar, - clear_calendars, ) -from zipline.utils.calendars.trading_calendar import days_at_time +from zipline.utils.calendars.calendar_utils import register_calendar_type +from zipline.utils.calendars.trading_calendar import days_at_time, \ + TradingCalendar + + +class FakeCalendar(TradingCalendar): + @property + def name(self): + return "DMY" + + @property + def tz(self): + return "Asia/Ulaanbaatar" + + @property + def open_time(self): + return time(11, 13) + + @property + def close_time(self): + return time(11, 49) class CalendarRegistrationTestCase(TestCase): - def setUp(self): - self.dummy_cal_type = namedtuple('DummyCal', ('name')) + self.dummy_cal_type = FakeCalendar def tearDown(self): - clear_calendars() + deregister_calendar('DMY') def test_register_calendar(self): # Build a fake calendar - dummy_cal = self.dummy_cal_type('DMY') + dummy_cal = self.dummy_cal_type() # Try to register and retrieve the calendar register_calendar('DMY', dummy_cal) @@ -69,18 +86,37 @@ class CalendarRegistrationTestCase(TestCase): with self.assertRaises(InvalidCalendarName): get_calendar('DMY') + def test_register_calendar_type(self): + register_calendar_type("DMY", self.dummy_cal_type) + retr_cal = get_calendar("DMY") + self.assertEqual(self.dummy_cal_type, type(retr_cal)) + + def test_both_places_are_checked(self): + dummy_cal = self.dummy_cal_type() + + # if instance is registered, can't register type with same name + register_calendar('DMY', dummy_cal) + with self.assertRaises(CalendarNameCollision): + register_calendar_type('DMY', type(dummy_cal)) + + deregister_calendar('DMY') + + # if type is registered, can't register instance with same name + register_calendar_type('DMY', type(dummy_cal)) + + with self.assertRaises(CalendarNameCollision): + register_calendar('DMY', dummy_cal) + def test_force_registration(self): - dummy_nyse = self.dummy_cal_type('NYSE') + register_calendar("DMY", self.dummy_cal_type()) + first_dummy = get_calendar("DMY") - # Get the actual NYSE calendar - real_nyse = get_calendar('NYSE') + # force-register a new instance + register_calendar("DMY", self.dummy_cal_type(), force=True) - # Force a registration of the dummy NYSE - register_calendar("NYSE", dummy_nyse, force=True) + second_dummy = get_calendar("DMY") - # Ensure that the dummy overwrote the real calendar - retr_cal = get_calendar('NYSE') - self.assertNotEqual(real_nyse, retr_cal) + self.assertNotEqual(first_dummy, second_dummy) class DaysAtTimeTestCase(TestCase): diff --git a/zipline/utils/calendars/__init__.py b/zipline/utils/calendars/__init__.py index 49bcc4dc..9314fba1 100644 --- a/zipline/utils/calendars/__init__.py +++ b/zipline/utils/calendars/__init__.py @@ -17,9 +17,10 @@ from .trading_calendar import TradingCalendar from .calendar_utils import ( get_calendar, register_calendar, + register_calendar_type, deregister_calendar, clear_calendars ) __all__ = ['get_calendar', 'TradingCalendar', 'register_calendar', - 'deregister_calendar', 'clear_calendars'] + 'register_calendar_type', 'deregister_calendar', 'clear_calendars'] diff --git a/zipline/utils/calendars/calendar_utils.py b/zipline/utils/calendars/calendar_utils.py index dca52f4e..9e60ab60 100644 --- a/zipline/utils/calendars/calendar_utils.py +++ b/zipline/utils/calendars/calendar_utils.py @@ -101,16 +101,41 @@ class TradingCalendarDispatcher(object): CalendarNameCollision If a calendar is already registered with the given calendar's name. """ - # If we are forcing the registration, remove an existing calendar with - # the same name. if force: self.deregister_calendar(name) - if name in self._calendars: + if name in self._calendars or name in self._calendar_factories: raise CalendarNameCollision(calendar_name=name) self._calendars[name] = calendar + def register_calendar_type(self, name, calendar_type, force=False): + """ + Registers a calendar by type. + + Parameters + ---------- + name: str + The key with which to register this calendar. + calendar_type: type + The type of the calendar to register. + force : bool, optional + If True, old calendars will be overwritten on a name collision. + If False, name collisions will raise an exception. Default: False. + + Raises + ------ + CalendarNameCollision + If a calendar is already registered with the given calendar's name. + """ + if force: + self._calendar_factories.pop(name, None) + + if name in self._calendars or name in self._calendar_factories: + raise CalendarNameCollision(calendar_name=name) + + self._calendar_factories[name] = calendar_type + def deregister_calendar(self, name): """ If a calendar is registered with the given name, it is de-registered. @@ -121,12 +146,14 @@ class TradingCalendarDispatcher(object): The name of the calendar to be deregistered. """ self._calendars.pop(name, None) + self._calendar_factories.pop(name, None) def clear_calendars(self): """ Deregisters all current registered calendars """ self._calendars.clear() + self._calendar_factories.clear() # We maintain a global calendar dispatcher so that users can just do @@ -140,3 +167,4 @@ get_calendar = global_calendar_dispatcher.get_calendar clear_calendars = global_calendar_dispatcher.clear_calendars deregister_calendar = global_calendar_dispatcher.deregister_calendar register_calendar = global_calendar_dispatcher.register_calendar +register_calendar_type = global_calendar_dispatcher.register_calendar_type