From e8185a1512e5840b935ffa80cacf0e8dd6e3d280 Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Wed, 23 Mar 2016 15:41:22 -0400 Subject: [PATCH] MAINT: reorganize - move testing mixin to fixtures BUG: correctly create asset finder MAINT: rename fixture STY: fixes for flake8 STY: add space around assignment MAINT: add var back to constructor MAINT: remove unused import MAINT: compare var with None directly MAINT: fix merge errors --- tests/pipeline/base.py | 113 +-------------- tests/pipeline/test_buyback_auth.py | 29 ++-- tests/pipeline/test_dividends.py | 60 ++++---- tests/pipeline/test_earnings.py | 22 ++- zipline/pipeline/data/dividends.py | 8 +- zipline/pipeline/factors/events.py | 4 +- .../pipeline/loaders/blaze/buyback_auth.py | 38 +---- zipline/pipeline/loaders/blaze/dividends.py | 48 +------ zipline/pipeline/loaders/blaze/earnings.py | 16 +-- zipline/pipeline/loaders/blaze/events.py | 6 +- zipline/pipeline/loaders/dividends.py | 16 +-- zipline/pipeline/loaders/utils.py | 5 +- zipline/testing/fixtures.py | 136 ++++++++++++++++-- 13 files changed, 205 insertions(+), 296 deletions(-) diff --git a/tests/pipeline/base.py b/tests/pipeline/base.py index 9c1952d5..1f89e56f 100644 --- a/tests/pipeline/base.py +++ b/tests/pipeline/base.py @@ -1,34 +1,25 @@ """ Base class for Pipeline API unittests. """ -import abc from functools import wraps from unittest import TestCase -from nose_parameterized import parameterized import numpy as np from numpy import arange, prod -import pandas as pd from pandas import date_range, Int64Index, DataFrame -from pandas.util.testing import assert_series_equal from six import iteritems -from zipline.pipeline import Pipeline, TermGraph +from zipline.pipeline import TermGraph from zipline.pipeline.engine import SimplePipelineEngine from zipline.pipeline.term import AssetExists from zipline.testing import ( check_arrays, ExplodingObject, - gen_calendars, make_simple_equity_info, tmp_asset_finder, ) from zipline.utils.functional import dzip_exact -from zipline.utils.numpy_utils import ( - NaTD, - make_datetime64D -) from zipline.utils.pandas_utils import explode from zipline.utils.tradingcalendar import trading_day @@ -175,105 +166,3 @@ class BasePipelineTestCase(TestCase): @with_default_shape def ones_mask(self, shape): return np.ones(shape, dtype=bool) - - -class EventLoaderCommonMixin(object): - @abc.abstractproperty - def get_sids(cls): - raise NotImplementedError('get_sids') - - @abc.abstractproperty - def get_dataset(self): - raise NotImplementedError('get_dataset') - - @abc.abstractproperty - def loader_type(self): - raise NotImplementedError('loader_type') - - def loader_args(self, dates): - """Construct the base object to pass to the loader. - - Parameters - ---------- - dates : pd.DatetimeIndex - The dates we can serve. - - Returns - ------- - args : tuple[any] - The arguments to forward to the loader positionally. - """ - return dates, self.get_dataset() - - def setup_engine(self, dates): - """ - Make a Pipeline Enigne object based on the given dates. - """ - loader = self.loader_type(*self.loader_args(dates)) - return SimplePipelineEngine(lambda _: loader, dates, self.asset_finder) - - @staticmethod - def _compute_busday_offsets(announcement_dates): - """ - Compute expected business day offsets from a DataFrame of announcement - dates. - """ - # Column-vector of dates on which factor `compute` will be called. - raw_call_dates = announcement_dates.index.values.astype( - 'datetime64[D]' - )[:, None] - - # 2D array of dates containining expected nexg announcement. - raw_announce_dates = ( - announcement_dates.values.astype('datetime64[D]') - ) - - # Set NaTs to 0 temporarily because busday_count doesn't support NaT. - # We fill these entries with NaNs later. - whereNaT = raw_announce_dates == NaTD - raw_announce_dates[whereNaT] = make_datetime64D(0) - - # The abs call here makes it so that we can use this function to - # compute offsets for both next and previous earnings (previous - # earnings offsets come back negative). - expected = abs(np.busday_count( - raw_call_dates, - raw_announce_dates - ).astype(float)) - - expected[whereNaT] = np.nan - return pd.DataFrame( - data=expected, - columns=announcement_dates.columns, - index=announcement_dates.index, - ) - - @parameterized.expand(gen_calendars( - '2014-01-01', - '2014-01-31', - critical_dates=pd.to_datetime([ - '2014-01-05', - '2014-01-10', - '2014-01-15', - '2014-01-20', - ], utc=True), - )) - def test_compute(self, dates): - engine = self.setup_engine(dates) - cols = self.setup(dates) - - pipe = Pipeline( - columns=self.pipeline_columns - ) - - result = engine.run_pipeline( - pipe, - start_date=dates[0], - end_date=dates[-1], - ) - - for sid in self.get_sids(): - for col_name in cols.keys(): - assert_series_equal(result[col_name].xs(sid, level=1), - cols[col_name][sid], - check_names=False) diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index a854d90a..e86abda5 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -5,7 +5,6 @@ import blaze as bz from blaze.compute.core import swap_resources_into_scope import pandas as pd from six import iteritems -from .base import EventLoaderCommonMixin from zipline.pipeline.common import( BUYBACK_ANNOUNCEMENT_FIELD_NAME, @@ -39,7 +38,9 @@ from zipline.pipeline.loaders.utils import ( zip_with_floats, zip_with_dates ) -from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase +from zipline.testing.fixtures import ( + WithPipelineEventDataLoader, ZiplineTestCase +) date_intervals = [[None, '2014-01-04'], ['2014-01-05', '2014-01-09'], ['2014-01-10', None]] @@ -74,8 +75,8 @@ def get_expected_previous_values(zip_date_index_with_vals, }, index=dates) -class CashBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase, - EventLoaderCommonMixin): +class CashBuybackAuthLoaderTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Test for cash buyback authorizations dataset. """ @@ -118,8 +119,8 @@ class CashBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase, return cols -class ShareBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase, - EventLoaderCommonMixin): +class ShareBuybackAuthLoaderTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Test for share buyback authorizations dataset. """ @@ -167,11 +168,11 @@ class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): """ loader_type = BlazeCashBuybackAuthorizationsLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeCashBuybackAuthLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.data(pd.concat( pd.DataFrame({ BUYBACK_ANNOUNCEMENT_FIELD_NAME: @@ -191,11 +192,11 @@ class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase): """ loader_type = BlazeShareBuybackAuthorizationsLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeShareBuybackAuthLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.data(pd.concat( pd.DataFrame({ BUYBACK_ANNOUNCEMENT_FIELD_NAME: @@ -214,11 +215,11 @@ class BlazeShareBuybackAuthLoaderNotInteractiveTestCase( BlazeShareBuybackAuthLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeShareBuybackAuthLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) @@ -226,9 +227,9 @@ class BlazeCashBuybackAuthLoaderNotInteractiveTestCase( BlazeCashBuybackAuthLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeCashBuybackAuthLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) diff --git a/tests/pipeline/test_dividends.py b/tests/pipeline/test_dividends.py index dc3ed210..431078b3 100644 --- a/tests/pipeline/test_dividends.py +++ b/tests/pipeline/test_dividends.py @@ -5,7 +5,6 @@ import blaze as bz from blaze.compute.core import swap_resources_into_scope import pandas as pd from six import iteritems -from tests.pipeline.base import EventLoaderCommonMixin from zipline.pipeline.common import ( ANNOUNCEMENT_FIELD_NAME, @@ -50,7 +49,10 @@ from zipline.pipeline.loaders.utils import ( zip_with_dates, zip_with_floats ) -from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase +from zipline.testing.fixtures import ( + WithPipelineEventDataLoader, + ZiplineTestCase +) dividends_cases = [ # K1--K2--A1--A2. @@ -184,8 +186,8 @@ def get_vals_for_dates(zip_date_index_with_vals, }, index=dates) -class DividendsByAnnouncementDateTestCase(WithAssetFinder, ZiplineTestCase, - EventLoaderCommonMixin): +class DividendsByAnnouncementDateTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Tests for loading the dividends by announcement date data. """ @@ -197,10 +199,6 @@ class DividendsByAnnouncementDateTestCase(WithAssetFinder, ZiplineTestCase, BusinessDaysSinceDividendAnnouncement(), } - @classmethod - def get_sids(cls): - return range(0, 5) - @classmethod def get_dataset(cls): return {sid: @@ -254,11 +252,11 @@ class BlazeDividendsByAnnouncementDateTestCase( ): loader_type = BlazeDividendsByAnnouncementDateLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeDividendsByAnnouncementDateTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.Data(pd.concat( pd.DataFrame({ ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME], @@ -275,32 +273,27 @@ class BlazeDividendsByAnnouncementDateNotInteractiveTestCase( """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeDividendsByAnnouncementDateNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) -class DividendsByExDateTestCase(WithAssetFinder, ZiplineTestCase, - EventLoaderCommonMixin): +class DividendsByExDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase): """ Tests for loading the dividends by ex date data. """ pipeline_columns = { - NEXT_EX_DATE: DividendsByExDate.next_ex_date.latest, - PREVIOUS_EX_DATE: DividendsByExDate.previous_ex_date.latest, + NEXT_EX_DATE: DividendsByExDate.next_date.latest, + PREVIOUS_EX_DATE: DividendsByExDate.previous_date.latest, NEXT_AMOUNT: DividendsByExDate.next_amount.latest, PREVIOUS_AMOUNT: DividendsByExDate.previous_amount.latest, DAYS_TO_NEXT_EX_DATE: BusinessDaysUntilNextExDate(), DAYS_SINCE_PREV_EX_DATE: BusinessDaysSincePreviousExDate() } - @classmethod - def get_sids(cls): - return range(0, 5) - @classmethod def get_dataset(cls): return {sid: @@ -342,11 +335,11 @@ class DividendsByExDateTestCase(WithAssetFinder, ZiplineTestCase, class BlazeDividendsByExDateLoaderTestCase(DividendsByExDateTestCase): loader_type = BlazeDividendsByExDateLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeDividendsByExDateLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.Data(pd.concat( pd.DataFrame({ EX_DATE_FIELD_NAME: df[EX_DATE_FIELD_NAME], @@ -363,30 +356,25 @@ class BlazeDividendsByExDateLoaderNotInteractiveTestCase( """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeDividendsByExDateLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) -class DividendsByPayDateTestCase(WithAssetFinder, ZiplineTestCase, - EventLoaderCommonMixin): +class DividendsByPayDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase): """ Tests for loading the dividends by pay date data. """ pipeline_columns = { - NEXT_PAY_DATE: DividendsByPayDate.next_pay_date.latest, - PREVIOUS_PAY_DATE: DividendsByPayDate.previous_pay_date.latest, + NEXT_PAY_DATE: DividendsByPayDate.next_date.latest, + PREVIOUS_PAY_DATE: DividendsByPayDate.previous_date.latest, NEXT_AMOUNT: DividendsByPayDate.next_amount.latest, PREVIOUS_AMOUNT: DividendsByPayDate.previous_amount.latest, } - @classmethod - def get_sids(cls): - return range(0, 5) - @classmethod def get_dataset(cls): return {sid: @@ -419,11 +407,11 @@ class DividendsByPayDateTestCase(WithAssetFinder, ZiplineTestCase, class BlazeDividendsByPayDateLoaderTestCase(DividendsByPayDateTestCase): loader_type = BlazeDividendsByPayDateLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeDividendsByPayDateLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.Data(pd.concat( pd.DataFrame({ PAY_DATE_FIELD_NAME: df[PAY_DATE_FIELD_NAME], @@ -440,9 +428,9 @@ class BlazeDividendsByPayDateLoaderNotInteractiveTestCase( """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeDividendsByPayDateLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index 0922a0c9..43f502c2 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -5,7 +5,6 @@ import blaze as bz from blaze.compute.core import swap_resources_into_scope import pandas as pd from six import iteritems -from .base import EventLoaderCommonMixin from zipline.pipeline.common import ( ANNOUNCEMENT_FIELD_NAME, @@ -28,7 +27,10 @@ from zipline.pipeline.loaders.utils import ( zip_with_dates ) -from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase +from zipline.testing.fixtures import ( + WithPipelineEventDataLoader, + ZiplineTestCase +) earnings_cases = [ # K1--K2--A1--A2. @@ -111,8 +113,8 @@ prev_dates = [ ] -class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase, - EventLoaderCommonMixin): +class EarningsCalendarLoaderTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Tests for loading the earnings announcement data. """ @@ -123,10 +125,6 @@ class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase, DAYS_TO_NEXT: BusinessDaysUntilNextEarnings(), } - @classmethod - def get_sids(cls): - return range(5) - @classmethod def get_dataset(cls): return {sid: df for sid, df in enumerate(earnings_cases)} @@ -199,11 +197,11 @@ class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase, class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase): loader_type = BlazeEarningsCalendarLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeEarningsCalendarLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.data(pd.concat( pd.DataFrame({ ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME], @@ -219,9 +217,9 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase( """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeEarningsCalendarLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) diff --git a/zipline/pipeline/data/dividends.py b/zipline/pipeline/data/dividends.py index b959e2a9..176aac43 100644 --- a/zipline/pipeline/data/dividends.py +++ b/zipline/pipeline/data/dividends.py @@ -7,15 +7,15 @@ from .dataset import Column, DataSet class DividendsByExDate(DataSet): - next_ex_date = Column(datetime64ns_dtype) - previous_ex_date = Column(datetime64ns_dtype) + next_date = Column(datetime64ns_dtype) + previous_date = Column(datetime64ns_dtype) next_amount = Column(float64_dtype) previous_amount = Column(float64_dtype) class DividendsByPayDate(DataSet): - next_pay_date = Column(datetime64ns_dtype) - previous_pay_date = Column(datetime64ns_dtype) + next_date = Column(datetime64ns_dtype) + previous_date = Column(datetime64ns_dtype) next_amount = Column(float64_dtype) previous_amount = Column(float64_dtype) diff --git a/zipline/pipeline/factors/events.py b/zipline/pipeline/factors/events.py index 100b7bfc..cdaa5573 100644 --- a/zipline/pipeline/factors/events.py +++ b/zipline/pipeline/factors/events.py @@ -189,7 +189,7 @@ class BusinessDaysUntilNextExDate( -------- zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement """ - inputs = [DividendsByExDate.next_ex_date] + inputs = [DividendsByExDate.next_date] class BusinessDaysSincePreviousExDate( @@ -204,4 +204,4 @@ class BusinessDaysSincePreviousExDate( -------- zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement """ - inputs = [DividendsByExDate.previous_ex_date] + inputs = [DividendsByExDate.previous_date] diff --git a/zipline/pipeline/loaders/blaze/buyback_auth.py b/zipline/pipeline/loaders/blaze/buyback_auth.py index 44b35c3e..9a1acfc9 100644 --- a/zipline/pipeline/loaders/blaze/buyback_auth.py +++ b/zipline/pipeline/loaders/blaze/buyback_auth.py @@ -68,24 +68,7 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader): }) concrete_loader = CashBuybackAuthorizationsLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=CashBuybackAuthorizations, - **kwargs): - super( - BlazeCashBuybackAuthorizationsLoader, self - ).__init__(expr, - resources=resources, - odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, - dataset=dataset, - **kwargs) + default_dataset = CashBuybackAuthorizations class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): @@ -140,21 +123,4 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): }) concrete_loader = ShareBuybackAuthorizationsLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=ShareBuybackAuthorizations, - **kwargs): - super( - BlazeShareBuybackAuthorizationsLoader, self - ).__init__(expr, - resources=resources, - odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, - dataset=dataset, - **kwargs) + default_dataset = ShareBuybackAuthorizations diff --git a/zipline/pipeline/loaders/blaze/dividends.py b/zipline/pipeline/loaders/blaze/dividends.py index 86a332b9..7bba9e05 100644 --- a/zipline/pipeline/loaders/blaze/dividends.py +++ b/zipline/pipeline/loaders/blaze/dividends.py @@ -72,21 +72,7 @@ class BlazeDividendsByAnnouncementDateLoader(BlazeEventsLoader): }) concrete_loader = DividendsByAnnouncementDateLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=DividendsByAnnouncementDate, - **kwargs): - super( - BlazeDividendsByAnnouncementDateLoader, self - ).__init__(expr, dataset=dataset, - resources=resources, odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, **kwargs) + default_dataset = DividendsByAnnouncementDate class BlazeDividendsByExDateLoader(BlazeEventsLoader): @@ -142,21 +128,7 @@ class BlazeDividendsByExDateLoader(BlazeEventsLoader): }) concrete_loader = DividendsByExDateLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=DividendsByExDate, - **kwargs): - super( - BlazeDividendsByExDateLoader, self - ).__init__(expr, dataset=dataset, - resources=resources, odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, **kwargs) + default_dataset = DividendsByExDate class BlazeDividendsByPayDateLoader(BlazeEventsLoader): @@ -212,18 +184,4 @@ class BlazeDividendsByPayDateLoader(BlazeEventsLoader): }) concrete_loader = DividendsByPayDateLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=DividendsByPayDate, - **kwargs): - super( - BlazeDividendsByPayDateLoader, self - ).__init__(expr, dataset=dataset, - resources=resources, odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, **kwargs) + default_dataset = DividendsByPayDate diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 17e76c63..9d7e9b5c 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -58,18 +58,4 @@ class BlazeEarningsCalendarLoader(BlazeEventsLoader): }) concrete_loader = EarningsCalendarLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=EarningsCalendar, - **kwargs): - super( - BlazeEarningsCalendarLoader, self - ).__init__(expr, dataset=dataset, - resources=resources, odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, **kwargs) + default_dataset = EarningsCalendar diff --git a/zipline/pipeline/loaders/blaze/events.py b/zipline/pipeline/loaders/blaze/events.py index 326c97ad..6fae418f 100644 --- a/zipline/pipeline/loaders/blaze/events.py +++ b/zipline/pipeline/loaders/blaze/events.py @@ -56,6 +56,7 @@ class BlazeEventsLoader(PipelineLoader): If the '{TS_FIELD_NAME}' field is not included it is assumed that we start the backtest with knowledge of all announcements. """ + default_dataset = None @preprocess(data_query_tz=optionally(ensure_timezone)) def __init__(self, @@ -64,7 +65,10 @@ class BlazeEventsLoader(PipelineLoader): odo_kwargs=None, data_query_time=None, data_query_tz=None, - dataset=None): + dataset=default_dataset): + if dataset is None: + dataset = self.default_dataset + dshape = expr.dshape if not istabular(dshape): diff --git a/zipline/pipeline/loaders/dividends.py b/zipline/pipeline/loaders/dividends.py index 4f80bf65..16a4585e 100644 --- a/zipline/pipeline/loaders/dividends.py +++ b/zipline/pipeline/loaders/dividends.py @@ -52,14 +52,14 @@ class DividendsByPayDateLoader(EventsLoader): ) @lazyval - def next_pay_date_loader(self): - return self._next_event_date_loader(self.dataset.next_pay_date, + def next_date_loader(self): + return self._next_event_date_loader(self.dataset.next_date, PAY_DATE_FIELD_NAME) @lazyval - def previous_pay_date_loader(self): + def previous_date_loader(self): return self._previous_event_date_loader( - self.dataset.previous_pay_date, + self.dataset.previous_date, PAY_DATE_FIELD_NAME ) @@ -90,14 +90,14 @@ class DividendsByExDateLoader(EventsLoader): ) @lazyval - def next_ex_date_loader(self): - return self._next_event_date_loader(self.dataset.next_ex_date, + def next_date_loader(self): + return self._next_event_date_loader(self.dataset.next_date, EX_DATE_FIELD_NAME) @lazyval - def previous_ex_date_loader(self): + def previous_date_loader(self): return self._previous_event_date_loader( - self.dataset.previous_ex_date, + self.dataset.previous_date, EX_DATE_FIELD_NAME ) diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py index 76d9f43c..3b79253f 100644 --- a/zipline/pipeline/loaders/utils.py +++ b/zipline/pipeline/loaders/utils.py @@ -60,8 +60,9 @@ def next_event_frame(events_by_sid, # Iterate over the raw Series values, since we're comparing against # numpy arrays anyway. - iterkv = zip(event_dates.index.values, event_dates.values, values) - for knowledge_date, event_date, value in iterkv: + iter_date_vals = zip(event_dates.index.values, event_dates.values, + values) + for knowledge_date, event_date, value in iter_date_vals: date_mask = ( (knowledge_date <= raw_dates) & (raw_dates <= event_date) diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index e4caf7ae..0e792dce 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -2,13 +2,19 @@ from unittest import TestCase from contextlib2 import ExitStack from logbook import NullHandler +from nose_parameterized import parameterized +import numpy as np import pandas as pd +from pandas.util.testing import assert_series_equal from six import with_metaclass -from .core import tmp_asset_finder, make_simple_equity_info +from .core import tmp_asset_finder, make_simple_equity_info, gen_calendars from ..finance.trading import TradingEnvironment from ..utils import tradingcalendar, factory from ..utils.final import FinalMeta, final +from zipline.pipeline import Pipeline, SimplePipelineEngine +from zipline.utils.numpy_utils import make_datetime64D +from zipline.utils.numpy_utils import NaTD class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)): @@ -177,14 +183,7 @@ class WithAssetFinder(object): def _make_info(cls): return None - @classmethod - def make_equities_info(cls): - return make_simple_equity_info( - cls.get_sids(), - start_date=pd.Timestamp('2013-01-01', tz='UTC'), - end_date=pd.Timestamp('2015-01-01', tz='UTC'), - ) - + make_equities_info = _make_info make_futures_info = _make_info make_exchanges_info = _make_info make_root_symbols_info = _make_info @@ -299,3 +298,122 @@ class WithNYSETradingDays(object): start_loc = end_loc - cls.TRADING_DAY_COUNT cls.trading_days = all_days[start_loc:end_loc + 1] + + +class WithPipelineEventDataLoader(WithAssetFinder): + """ + ZiplineTestCase mixin providing common test methods/behaviors for event + data loaders. + + `get_sids` must return the sids being tested. + `get_dataset` must return {sid -> pd.DataFrame} + `loader_type` must return the loader class to use for loading the dataset + `make_asset_finder` returns a default asset finder which can be overridden. + """ + @classmethod + def get_sids(cls): + return range(0, 5) + + @classmethod + def get_dataset(cls): + return {sid: pd.DataFrame() for sid in cls.get_sids()} + + @classmethod + def loader_type(self): + return None + + @classmethod + def make_equities_info(cls): + return make_simple_equity_info( + cls.get_sids(), + start_date=pd.Timestamp('2013-01-01', tz='UTC'), + end_date=pd.Timestamp('2015-01-01', tz='UTC'), + ) + + def pipeline_event_loader_args(self, dates): + """Construct the base object to pass to the loader. + + Parameters + ---------- + dates : pd.DatetimeIndex + The dates we can serve. + + Returns + ------- + args : tuple[any] + The arguments to forward to the loader positionally. + """ + return dates, self.get_dataset() + + def pipeline_event_setup_engine(self, dates): + """ + Make a Pipeline Enigne object based on the given dates. + """ + loader = self.loader_type(*self.pipeline_event_loader_args(dates)) + return SimplePipelineEngine(lambda _: loader, dates, self.asset_finder) + + @staticmethod + def _compute_busday_offsets(announcement_dates): + """ + Compute expected business day offsets from a DataFrame of announcement + dates. + """ + # Column-vector of dates on which factor `compute` will be called. + raw_call_dates = announcement_dates.index.values.astype( + 'datetime64[D]' + )[:, None] + + # 2D array of dates containining expected nexg announcement. + raw_announce_dates = ( + announcement_dates.values.astype('datetime64[D]') + ) + + # Set NaTs to 0 temporarily because busday_count doesn't support NaT. + # We fill these entries with NaNs later. + whereNaT = raw_announce_dates == NaTD + raw_announce_dates[whereNaT] = make_datetime64D(0) + + # The abs call here makes it so that we can use this function to + # compute offsets for both next and previous earnings (previous + # earnings offsets come back negative). + expected = abs(np.busday_count( + raw_call_dates, + raw_announce_dates + ).astype(float)) + + expected[whereNaT] = np.nan + return pd.DataFrame( + data=expected, + columns=announcement_dates.columns, + index=announcement_dates.index, + ) + + @parameterized.expand(gen_calendars( + '2014-01-01', + '2014-01-31', + critical_dates=pd.to_datetime([ + '2014-01-05', + '2014-01-10', + '2014-01-15', + '2014-01-20', + ], utc=True), + )) + def test_compute(self, dates): + engine = self.pipeline_event_setup_engine(dates) + cols = self.setup(dates) + + pipe = Pipeline( + columns=self.pipeline_columns + ) + + result = engine.run_pipeline( + pipe, + start_date=dates[0], + end_date=dates[-1], + ) + + for sid in self.get_sids(): + for col_name in cols.keys(): + assert_series_equal(result[col_name].xs(sid, level=1), + cols[col_name][sid], + check_names=False)