From 2a81c2066f360f0578120ef093c1925953fe79db Mon Sep 17 00:00:00 2001
From: jfkirk <james.f.kirk@gmail.com>
Date: Wed, 25 May 2016 11:16:25 -0400
Subject: [PATCH] ENH: Adds the option to force calendar registration

---
 tests/test_exchange_calendar.py              | 51 ++++++++++++++++++++
 zipline/utils/calendars/exchange_calendar.py | 46 +++++++++++++++++-
 2 files changed, 96 insertions(+), 1 deletion(-)

diff --git a/tests/test_exchange_calendar.py b/tests/test_exchange_calendar.py
index d21a1f29..cb75db97 100644
--- a/tests/test_exchange_calendar.py
+++ b/tests/test_exchange_calendar.py
@@ -19,6 +19,7 @@ from os.path import (
     join,
 )
 from unittest import TestCase
+from collections import namedtuple
 
 import pandas as pd
 import pytz
@@ -31,7 +32,57 @@ from pandas import (
 )
 from pandas.util.testing import assert_frame_equal
 
+from zipline.errors import (
+    CalendarNameCollision,
+    InvalidCalendarName,
+)
 from zipline.utils.calendars.exchange_calendar_nyse import NYSEExchangeCalendar
+from zipline.utils.calendars.exchange_calendar import(
+    register_calendar,
+    deregister_calendar,
+    get_calendar,
+    clear_calendars,
+)
+
+
+class CalendarRegistrationTestCase(TestCase):
+
+    def setUp(self):
+        self.dummy_cal_type = namedtuple('DummyCal', ('name'))
+
+    def tearDown(self):
+        clear_calendars()
+
+    def test_register_calendar(self):
+        # Build a fake calendar
+        dummy_cal = self.dummy_cal_type('DMY')
+
+        # Try to register and retrieve the calendar
+        register_calendar(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)
+
+        # Deregister the calendar and ensure that it is removed
+        deregister_calendar('DMY')
+        with self.assertRaises(InvalidCalendarName):
+            get_calendar('DMY')
+
+    def test_force_registration(self):
+        dummy_nyse = self.dummy_cal_type('NYSE')
+
+        # Get the actual NYSE calendar
+        real_nyse = get_calendar('NYSE')
+
+        # Force a registration of the dummy NYSE
+        register_calendar(dummy_nyse, force=True)
+
+        # Ensure that the dummy overwrote the real calendar
+        retr_cal = get_calendar('NYSE')
+        self.assertNotEqual(real_nyse, retr_cal)
 
 
 class ExchangeCalendarTestBase(object):
diff --git a/zipline/utils/calendars/exchange_calendar.py b/zipline/utils/calendars/exchange_calendar.py
index 6bf66e78..495e97fb 100644
--- a/zipline/utils/calendars/exchange_calendar.py
+++ b/zipline/utils/calendars/exchange_calendar.py
@@ -496,8 +496,52 @@ def get_calendar(name):
     return _static_calendars[name]
 
 
-def register_calendar(calendar):
+def deregister_calendar(cal_name):
+    """
+    If a calendar is registered with the given name, it is de-registered.
+
+    Parameters
+    ----------
+    cal_name : str
+        The name of the calendar to be deregistered.
+    """
+    try:
+        _static_calendars.pop(cal_name)
+    except KeyError:
+        pass
+
+
+def clear_calendars():
+    """
+    Deregisters all current registered calendars
+    """
+    _static_calendars.clear()
+
+
+def register_calendar(calendar, force=False):
+    """
+    Registers a calendar for retrieval by the get_calendar method.
+
+    Parameters
+    ----------
+    calendar : ExchangeCalendar
+        The calendar to be registered for retrieval.
+    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 we are forcing the registration, remove an existing calendar with the
+    # same name.
+    if force:
+        deregister_calendar(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)
+
     _static_calendars[calendar.name] = calendar