diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 6c90b16b..6b5c1695 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/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) diff --git a/zipline/pipeline/loaders/earnings.py b/zipline/pipeline/loaders/earnings.py index cb8106b6..b5792a1d 100644 --- a/zipline/pipeline/loaders/earnings.py +++ b/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,