ENH: support subclassing in the earningscalendar loader

zipline/pipeline/loaders/blaze/earnings.py

@@ -6,6 +6,7 @@ from six import iteritems
 from toolz import valmap
 
 from .core import TS_FIELD_NAME, SID_FIELD_NAME
+from zipline.pipeline.data import EarningsCalendar
 from zipline.pipeline.loaders.base import PipelineLoader
 from zipline.pipeline.loaders.earnings import EarningsCalendarLoader
 
@@ -87,7 +88,8 @@ class BlazeEarningsCalendarLoader(PipelineLoader):
                  expr,
                  resources=None,
                  compute_kwargs=None,
-                 odo_kwargs=None):
+                 odo_kwargs=None,
+                 dataset=EarningsCalendar):
         dshape = expr.dshape
 
         if not istabular(dshape):
@@ -101,6 +103,7 @@ class BlazeEarningsCalendarLoader(PipelineLoader):
             resources,
         )
         self._odo_kwargs = odo_kwargs if odo_kwargs is not None else {}
+        self._dataset = dataset
 
     def load_adjusted_array(self, columns, dates, assets, mask):
         expr = self._expr
@@ -147,4 +150,5 @@ class BlazeEarningsCalendarLoader(PipelineLoader):
         return EarningsCalendarLoader(
             dates,
             valmap(mkseries, gb.groups),
+            dataset=self._dataset,
         ).load_adjusted_array(columns, dates, assets, mask)

zipline/pipeline/loaders/earnings.py

@@ -44,9 +44,12 @@ class EarningsCalendarLoader(PipelineLoader):
         Whether to allow passing ``DatetimeIndex`` values in
         ``announcement_dates``.
     """
-    def __init__(self, all_dates, announcement_dates, infer_timestamps=False):
+    def __init__(self,
+                 all_dates,
+                 announcement_dates,
+                 infer_timestamps=False,
+                 dataset=EarningsCalendar):
         self.all_dates = all_dates
-
         self.announcement_dates = announcement_dates = (
             announcement_dates.copy()
         )
@@ -64,13 +67,14 @@ class EarningsCalendarLoader(PipelineLoader):
                 announcement_dates[k] = pd.Series(
                     v, index=repeat(dates[0], len(v)),
                 )
+        self.dataset = dataset
 
     def get_loader(self, column):
         """Dispatch to the loader for ``column``.
         """
-        if column is EarningsCalendar.next_announcement:
+        if column is self.dataset.next_announcement:
             return self.next_announcement_loader
-        elif column is EarningsCalendar.previous_announcement:
+        elif column is self.dataset.previous_announcement:
             return self.previous_announcement_loader
         else:
             raise ValueError("Don't know how to load column '%s'." % column)
@@ -78,7 +82,7 @@ class EarningsCalendarLoader(PipelineLoader):
     @lazyval
     def next_announcement_loader(self):
         return DataFrameLoader(
-            EarningsCalendar.next_announcement,
+            self.dataset.next_announcement,
             next_earnings_date_frame(
                 self.all_dates,
                 self.announcement_dates,
@@ -89,7 +93,7 @@ class EarningsCalendarLoader(PipelineLoader):
     @lazyval
     def previous_announcement_loader(self):
         return DataFrameLoader(
-            EarningsCalendar.previous_announcement,
+            self.dataset.previous_announcement,
             previous_earnings_date_frame(
                 self.all_dates,
                 self.announcement_dates,