Merge pull request quantopian#1385 from quantopian/register-calendar-…

…types

ENH: Add public API to register calendars by type

tests/calendars/test_trading_calendar.py

@@ -19,7 +19,6 @@
     join,
 )
 from unittest import TestCase
-from collections import namedtuple
 
 import numpy as np
 import pandas as pd
@@ -38,22 +37,40 @@
     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 CalendarRegistrationTestCase(TestCase):
+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 @@ def test_register_calendar(self):
         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):

zipline/utils/calendars/__init__.py

@@ -17,9 +17,10 @@
 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']

zipline/utils/calendars/calendar_utils.py

@@ -101,16 +101,41 @@ def register_calendar(self, name, calendar, force=False):
         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 @@ def deregister_calendar(self, name):
             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 @@ def clear_calendars(self):
 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