From 06dd6e958de5518927a517dc930631a642afd17b Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Thu, 17 Mar 2016 14:40:21 -0400 Subject: [PATCH] TST: recfator tests to use fixtures MAINT: use np.array MAINT: return cols rather than modifying attribute --- tests/pipeline/base.py | 32 +- tests/pipeline/test_buyback_auth.py | 154 +++------ tests/pipeline/test_dividends.py | 360 ++++++++------------ tests/pipeline/test_earnings.py | 218 ++++++------ zipline/pipeline/loaders/blaze/dividends.py | 14 +- zipline/pipeline/loaders/utils.py | 62 +++- zipline/testing/fixtures.py | 11 +- 7 files changed, 371 insertions(+), 480 deletions(-) diff --git a/tests/pipeline/base.py b/tests/pipeline/base.py index e8261ccb..9c1952d5 100644 --- a/tests/pipeline/base.py +++ b/tests/pipeline/base.py @@ -21,7 +21,6 @@ from zipline.testing import ( ExplodingObject, gen_calendars, make_simple_equity_info, - num_days_in_range, tmp_asset_finder, ) @@ -183,22 +182,13 @@ class EventLoaderCommonMixin(object): def get_sids(cls): raise NotImplementedError('get_sids') - @classmethod - def get_equity_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'), - ) + @abc.abstractproperty + def get_dataset(self): + raise NotImplementedError('get_dataset') - def zip_with_floats(self, dates, flts): - return pd.Series(flts, index=dates).astype('float') - - def num_days_between(self, dates, start_date, end_date): - return num_days_in_range(dates, start_date, end_date) - - def zip_with_dates(self, index_dates, dts): - return pd.Series(pd.to_datetime(dts), index=index_dates) + @abc.abstractproperty + def loader_type(self): + raise NotImplementedError('loader_type') def loader_args(self, dates): """Construct the base object to pass to the loader. @@ -213,14 +203,14 @@ class EventLoaderCommonMixin(object): args : tuple[any] The arguments to forward to the loader positionally. """ - return dates, self.dataset + 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.finder) + return SimplePipelineEngine(lambda _: loader, dates, self.asset_finder) @staticmethod def _compute_busday_offsets(announcement_dates): @@ -270,7 +260,7 @@ class EventLoaderCommonMixin(object): )) def test_compute(self, dates): engine = self.setup_engine(dates) - self.setup(dates) + cols = self.setup(dates) pipe = Pipeline( columns=self.pipeline_columns @@ -283,7 +273,7 @@ class EventLoaderCommonMixin(object): ) for sid in self.get_sids(): - for col_name in self.cols.keys(): + for col_name in cols.keys(): assert_series_equal(result[col_name].xs(sid, level=1), - self.cols[col_name][sid], + cols[col_name][sid], check_names=False) diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index 35fe592f..a854d90a 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -1,13 +1,8 @@ """ Tests for the reference loader for Buyback Authorizations. """ -from functools import partial -from unittest import TestCase - import blaze as bz from blaze.compute.core import swap_resources_into_scope -from contextlib2 import ExitStack -import itertools import pandas as pd from six import iteritems from .base import EventLoaderCommonMixin @@ -39,7 +34,12 @@ from zipline.pipeline.loaders.blaze import ( BlazeCashBuybackAuthorizationsLoader, BlazeShareBuybackAuthorizationsLoader, ) -from zipline.testing import tmp_asset_finder +from zipline.pipeline.loaders.utils import ( + get_values_for_date_ranges, + zip_with_floats, + zip_with_dates +) +from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase date_intervals = [[None, '2014-01-04'], ['2014-01-05', '2014-01-09'], ['2014-01-10', None]] @@ -62,48 +62,20 @@ buyback_authorizations_cases = [ ] -def get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals_for_date_intervals): - # Fill in given values for given date ranges. - return zip_with_floats_dates( - list( - itertools.chain(*[ - [val] * num_days_between_dates(*date_intervals[i]) - for i, val in enumerate(vals_for_date_intervals) - ]) - ) - ) - - -def get_expected_previous_values(zip_with_floats_dates, - num_days_between_dates, +def get_expected_previous_values(zip_date_index_with_vals, dates, vals_for_date_intervals): return pd.DataFrame({ - 0: get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals_for_date_intervals), - 1: zip_with_floats_dates(['NaN'] * len(dates)), + 0: get_values_for_date_ranges(zip_date_index_with_vals, + vals_for_date_intervals, + date_intervals, + dates), + 1: zip_date_index_with_vals(dates, ['NaN'] * len(dates)), }, index=dates) -def get_expected_previous_dates(zip_with_dates_for_dates, - num_days_between_for_dates, - dates): - return pd.DataFrame({ - 0: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-04'] * num_days_between_for_dates('2014-01-05', - '2014-01-09') + - ['2014-01-09'] * num_days_between_for_dates('2014-01-10', - None), - ), - 1: zip_with_dates_for_dates(['NaT'] * len(dates)) - }) - - -class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): +class CashBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase, + EventLoaderCommonMixin): """ Test for cash buyback authorizations dataset. """ @@ -121,43 +93,33 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): return range(2) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - cls.finder = stack.enter_context( - tmp_asset_finder(equities=cls.get_equity_info()), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop(SHARE_COUNT_FIELD_NAME, axis=1) - for sid, frame - in enumerate(buyback_authorizations_cases)} - cls.loader_type = CashBuybackAuthorizationsLoader + def get_dataset(cls): + return {sid: + frame.drop(SHARE_COUNT_FIELD_NAME, axis=1) + for sid, frame + in enumerate(buyback_authorizations_cases)} - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = CashBuybackAuthorizationsLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) + cols = {} _expected_previous_cash = get_expected_previous_values( - zip_with_floats_dates, num_days_between_dates, dates, + zip_with_floats, dates, ['NaN', 10, 20] ) - self.cols[ + cols[ PREVIOUS_BUYBACK_ANNOUNCEMENT - ] = get_expected_previous_dates(zip_with_dates_for_dates, - num_days_between_for_dates, - dates) - self.cols[PREVIOUS_BUYBACK_CASH] = _expected_previous_cash - self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( - self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] + ] = get_expected_previous_values(zip_with_dates, dates, + ['NaT', '2014-01-04', '2014-01-09']) + cols[PREVIOUS_BUYBACK_CASH] = _expected_previous_cash + cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( + cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] ) + return cols -class ShareBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): +class ShareBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase, + EventLoaderCommonMixin): """ Test for share buyback authorizations dataset. """ @@ -175,50 +137,35 @@ class ShareBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): return range(2) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - cls.finder = stack.enter_context( - tmp_asset_finder(equities=cls.get_equity_info()), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop(CASH_FIELD_NAME, axis=1) - for sid, frame - in enumerate(buyback_authorizations_cases)} - cls.loader_type = ShareBuybackAuthorizationsLoader + def get_dataset(cls): + return {sid: + frame.drop(CASH_FIELD_NAME, axis=1) + for sid, frame + in enumerate(buyback_authorizations_cases)} - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = ShareBuybackAuthorizationsLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) - - self.cols[ + cols = {} + cols[ PREVIOUS_BUYBACK_SHARE_COUNT - ] = get_expected_previous_values(zip_with_floats_dates, - num_days_between_dates, dates, + ] = get_expected_previous_values(zip_with_floats, + dates, ['NaN', 1, 15]) - self.cols[ + cols[ PREVIOUS_BUYBACK_ANNOUNCEMENT - ] = get_expected_previous_dates(zip_with_dates_for_dates, - num_days_between_for_dates, - dates) - self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( - self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] + ] = get_expected_previous_values(zip_with_dates, dates, + ['NaT', '2014-01-04', '2014-01-09']) + cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( + cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] ) + return cols class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): """ Test case for loading via blaze. """ - @classmethod - def setUpClass(cls): - super(BlazeCashBuybackAuthLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeCashBuybackAuthorizationsLoader + loader_type = BlazeCashBuybackAuthorizationsLoader def loader_args(self, dates): _, mapping = super( @@ -242,10 +189,7 @@ class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase): """ Test case for loading via blaze. """ - @classmethod - def setUpClass(cls): - super(BlazeShareBuybackAuthLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeShareBuybackAuthorizationsLoader + loader_type = BlazeShareBuybackAuthorizationsLoader def loader_args(self, dates): _, mapping = super( diff --git a/tests/pipeline/test_dividends.py b/tests/pipeline/test_dividends.py index 099bf67e..dc3ed210 100644 --- a/tests/pipeline/test_dividends.py +++ b/tests/pipeline/test_dividends.py @@ -1,13 +1,8 @@ """ Tests for the reference loader for Dividends datasets. """ -from functools import partial -from unittest import TestCase - import blaze as bz from blaze.compute.core import swap_resources_into_scope -from contextlib2 import ExitStack -import itertools import pandas as pd from six import iteritems from tests.pipeline.base import EventLoaderCommonMixin @@ -30,23 +25,32 @@ from zipline.pipeline.common import ( EX_DATE_FIELD_NAME, PAY_DATE_FIELD_NAME ) -from zipline.pipeline.data.dividends import DividendsByAnnouncementDate, \ - DividendsByExDate, DividendsByPayDate +from zipline.pipeline.data.dividends import ( + DividendsByAnnouncementDate, + DividendsByExDate, + DividendsByPayDate +) from zipline.pipeline.factors.events import ( BusinessDaysSinceDividendAnnouncement, BusinessDaysSincePreviousExDate, BusinessDaysUntilNextExDate ) -from zipline.pipeline.loaders.blaze.dividends import \ - BlazeDividendsByAnnouncementDateLoader, BlazeDividendsByPayDateLoader, \ +from zipline.pipeline.loaders.blaze.dividends import ( + BlazeDividendsByAnnouncementDateLoader, + BlazeDividendsByPayDateLoader, BlazeDividendsByExDateLoader -from zipline.pipeline.loaders.dividends import DividendsByAnnouncementDateLoader, \ - DividendsByExDateLoader, DividendsByPayDateLoader -from zipline.utils.test_utils import ( - make_simple_equity_info, - tmp_asset_finder, ) - +from zipline.pipeline.loaders.dividends import ( + DividendsByAnnouncementDateLoader, + DividendsByExDateLoader, + DividendsByPayDateLoader +) +from zipline.pipeline.loaders.utils import ( + get_values_for_date_ranges, + zip_with_dates, + zip_with_floats +) +from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase dividends_cases = [ # K1--K2--A1--A2. @@ -92,23 +96,23 @@ dividends_cases = [ ] prev_date_intervals = [ - [ - [None, '2014-01-14'], ['2014-01-15', '2014-01-19'], - ['2014-01-20', None] - ], - [ - [None, '2014-01-14'], ['2014-01-15', '2014-01-19'], - ['2014-01-20', None] - ], - [ - [None, '2014-01-09'], ['2014-01-10', '2014-01-19'], - ['2014-01-20', None] - ], - [ - [None, '2014-01-09'], ['2014-01-10', '2014-01-14'], - ['2014-01-15', None] - ] - ] + [ + [None, '2014-01-14'], ['2014-01-15', '2014-01-19'], + ['2014-01-20', None] + ], + [ + [None, '2014-01-14'], ['2014-01-15', '2014-01-19'], + ['2014-01-20', None] + ], + [ + [None, '2014-01-09'], ['2014-01-10', '2014-01-19'], + ['2014-01-20', None] + ], + [ + [None, '2014-01-09'], ['2014-01-10', '2014-01-14'], + ['2014-01-15', None] + ] +] next_date_intervals = [ [ @@ -138,9 +142,9 @@ next_ex_and_pay_dates = [['NaT', '2014-01-15', '2014-01-20', 'NaT'], ['NaT', '2014-01-10', '2014-01-15', 'NaT']] prev_ex_and_pay_dates = [['NaT', '2014-01-15', '2014-01-20'], - ['NaT', '2014-01-15', '2014-01-20'], - ['NaT', '2014-01-10', '2014-01-20'], - ['NaT', '2014-01-10', '2014-01-15']] + ['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-15']] prev_amounts = [['NaN', 1, 15], ['NaN', 13, 7], @@ -153,50 +157,35 @@ next_amounts = [['NaN', 1, 15, 'NaN'], ['NaN', 6, 23, 'NaN']] -def get_values_for_date_ranges(zip_vals_dates, - num_days_between_dates, - vals_for_date_intervals, - date_intervals): - # Fill in given values for given date ranges. - return zip_vals_dates( - list( - itertools.chain(*[ - [val] * num_days_between_dates(*date_intervals[i]) - for i, val in enumerate(vals_for_date_intervals) - ]) - ) - ) - - -def get_vals_for_dates(zip_with_floats_dates, - num_days_between_dates, - dates, - date_invervals, - vals): +def get_vals_for_dates(zip_date_index_with_vals, + vals, + date_invervals, + dates): return pd.DataFrame({ - 0: get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals[0], - date_invervals[0]), - 1: get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals[1], - date_invervals[1]), - 2: get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals[2], - date_invervals[2]), - # Assume the latest of 2 cash values is used if we find out about 2 - # announcements that happened on the same day for the same sid. - 3: get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals[3], - date_invervals[3]), - 4: zip_with_floats_dates(['NaN'] * len(dates)), - }, index=dates) + 0: get_values_for_date_ranges(zip_date_index_with_vals, + vals[0], + date_invervals[0], + dates), + 1: get_values_for_date_ranges(zip_date_index_with_vals, + vals[1], + date_invervals[1], + dates), + 2: get_values_for_date_ranges(zip_date_index_with_vals, + vals[2], + date_invervals[2], + dates), + # Assume the latest of 2 cash values is used if we find out about 2 + # announcements that happened on the same day for the same sid. + 3: get_values_for_date_ranges(zip_date_index_with_vals, + vals[3], + date_invervals[3], + dates), + 4: zip_date_index_with_vals(dates, ['NaN'] * len(dates)), + }, index=dates) -class DividendsByAnnouncementDateTestCase(TestCase, EventLoaderCommonMixin): +class DividendsByAnnouncementDateTestCase(WithAssetFinder, ZiplineTestCase, + EventLoaderCommonMixin): """ Tests for loading the dividends by announcement date data. """ @@ -213,34 +202,16 @@ class DividendsByAnnouncementDateTestCase(TestCase, EventLoaderCommonMixin): return range(0, 5) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - equity_info = 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'), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop([EX_DATE_FIELD_NAME, - PAY_DATE_FIELD_NAME], axis=1) - for sid, frame - in enumerate(dividends_cases)} - cls.finder = stack.enter_context( - tmp_asset_finder(equities=equity_info), - ) + def get_dataset(cls): + return {sid: + frame.drop([EX_DATE_FIELD_NAME, + PAY_DATE_FIELD_NAME], axis=1) + for sid, frame + in enumerate(dividends_cases)} - cls.loader_type = DividendsByAnnouncementDateLoader - - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = DividendsByAnnouncementDateLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) date_intervals = [ [ [None, '2014-01-04'], ['2014-01-05', '2014-01-09'], @@ -251,41 +222,37 @@ class DividendsByAnnouncementDateTestCase(TestCase, EventLoaderCommonMixin): ['2014-01-10', None] ], [ - [None, '2014-01-04'], ['2014-01-05', '2014-01-14'], - ['2014-01-15', None] + [None, '2014-01-04'], ['2014-01-05', '2014-01-14'], + ['2014-01-15', None] ], [ - [None, '2014-01-04'], ['2014-01-05', None] + [None, '2014-01-04'], ['2014-01-05', None] ] ] announcement_dates = [['NaT', '2014-01-04', '2014-01-09'], - ['NaT', '2014-01-04', '2014-01-09'], - ['NaT', '2014-01-04', '2014-01-14'], - ['NaT', '2014-01-04']] + ['NaT', '2014-01-04', '2014-01-09'], + ['NaT', '2014-01-04', '2014-01-14'], + ['NaT', '2014-01-04']] amounts = [['NaN', 1, 15], ['NaN', 7, 13], ['NaN', 3, 1], ['NaN', 23]] - - self.cols[PREVIOUS_ANNOUNCEMENT] = get_vals_for_dates( - zip_with_dates_for_dates, num_days_between_for_dates, dates, - date_intervals, announcement_dates + cols = {} + cols[PREVIOUS_ANNOUNCEMENT] = get_vals_for_dates( + zip_with_dates, announcement_dates, date_intervals, dates ) - self.cols[PREVIOUS_AMOUNT] = get_vals_for_dates( - zip_with_floats_dates, num_days_between_dates, dates, - date_intervals, amounts + cols[PREVIOUS_AMOUNT] = get_vals_for_dates( + zip_with_floats, amounts, date_intervals, dates ) - self.cols[ + cols[ DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT - ] = self._compute_busday_offsets(self.cols[PREVIOUS_ANNOUNCEMENT]) + ] = self._compute_busday_offsets(cols[PREVIOUS_ANNOUNCEMENT]) + return cols class BlazeDividendsByAnnouncementDateTestCase( DividendsByAnnouncementDateTestCase ): - @classmethod - def setUpClass(cls): - super(BlazeDividendsByAnnouncementDateTestCase, cls).setUpClass() - cls.loader_type = BlazeDividendsByAnnouncementDateLoader + loader_type = BlazeDividendsByAnnouncementDateLoader def loader_args(self, dates): _, mapping = super( @@ -307,11 +274,6 @@ class BlazeDividendsByAnnouncementDateNotInteractiveTestCase( BlazeDividendsByAnnouncementDateTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - @classmethod - def setUpClass(cls): - super(BlazeDividendsByAnnouncementDateNotInteractiveTestCase, - cls).setUpClass() - cls.loader_type = BlazeDividendsByAnnouncementDateLoader def loader_args(self, dates): (bound_expr,) = super( @@ -321,13 +283,14 @@ class BlazeDividendsByAnnouncementDateNotInteractiveTestCase( return swap_resources_into_scope(bound_expr, {}) -class DividendsByExDateTestCase(TestCase, EventLoaderCommonMixin): +class DividendsByExDateTestCase(WithAssetFinder, ZiplineTestCase, + EventLoaderCommonMixin): """ Tests for loading the dividends by ex date data. """ pipeline_columns = { - NEXT_EX_DATE: DividendsByExDate.previous_ex_date.latest, - PREVIOUS_EX_DATE: DividendsByExDate.next_ex_date.latest, + NEXT_EX_DATE: DividendsByExDate.next_ex_date.latest, + PREVIOUS_EX_DATE: DividendsByExDate.previous_ex_date.latest, NEXT_AMOUNT: DividendsByExDate.next_amount.latest, PREVIOUS_AMOUNT: DividendsByExDate.previous_amount.latest, DAYS_TO_NEXT_EX_DATE: BusinessDaysUntilNextExDate(), @@ -339,69 +302,45 @@ class DividendsByExDateTestCase(TestCase, EventLoaderCommonMixin): return range(0, 5) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - equity_info = 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'), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop([ANNOUNCEMENT_FIELD_NAME, - PAY_DATE_FIELD_NAME], axis=1) - for sid, frame - in enumerate(dividends_cases)} - cls.finder = stack.enter_context( - tmp_asset_finder(equities=equity_info), - ) + def get_dataset(cls): + return {sid: + frame.drop([ANNOUNCEMENT_FIELD_NAME, + PAY_DATE_FIELD_NAME], axis=1) + for sid, frame + in enumerate(dividends_cases)} - cls.loader_type = DividendsByExDateLoader - - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = DividendsByExDateLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) - - self.cols[NEXT_EX_DATE] = get_vals_for_dates( - zip_with_dates_for_dates, num_days_between_for_dates, dates, - next_date_intervals, next_ex_and_pay_dates + cols = {} + cols[NEXT_EX_DATE] = get_vals_for_dates( + zip_with_dates, next_ex_and_pay_dates, next_date_intervals, dates, ) - self.cols[PREVIOUS_EX_DATE] = get_vals_for_dates( - zip_with_dates_for_dates, num_days_between_for_dates, dates, - prev_date_intervals, prev_ex_and_pay_dates + cols[PREVIOUS_EX_DATE] = get_vals_for_dates( + zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals, dates ) - self.cols[NEXT_AMOUNT] = get_vals_for_dates( - zip_with_floats_dates, num_days_between_dates, - dates, next_date_intervals, next_amounts + cols[NEXT_AMOUNT] = get_vals_for_dates( + zip_with_floats, next_amounts, next_date_intervals, dates ) - self.cols[PREVIOUS_AMOUNT] = get_vals_for_dates( - zip_with_floats_dates, num_days_between_dates, - dates, prev_date_intervals, prev_amounts + cols[PREVIOUS_AMOUNT] = get_vals_for_dates( + zip_with_floats, prev_amounts, prev_date_intervals, dates ) - self.cols[DAYS_TO_NEXT_EX_DATE] = self._compute_busday_offsets( - self.cols[NEXT_EX_DATE] + cols[DAYS_TO_NEXT_EX_DATE] = self._compute_busday_offsets( + cols[NEXT_EX_DATE] ) - self.cols[DAYS_SINCE_PREV_EX_DATE] = self._compute_busday_offsets( - self.cols[PREVIOUS_EX_DATE] + cols[DAYS_SINCE_PREV_EX_DATE] = self._compute_busday_offsets( + cols[PREVIOUS_EX_DATE] ) + return cols class BlazeDividendsByExDateLoaderTestCase(DividendsByExDateTestCase): - @classmethod - def setUpClass(cls): - super(BlazeDividendsByExDateLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeDividendsByExDateLoader + loader_type = BlazeDividendsByExDateLoader def loader_args(self, dates): _, mapping = super( @@ -423,11 +362,6 @@ class BlazeDividendsByExDateLoaderNotInteractiveTestCase( BlazeDividendsByExDateLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - @classmethod - def setUpClass(cls): - super(BlazeDividendsByExDateLoaderNotInteractiveTestCase, - cls).setUpClass() - cls.loader_type = DividendsByExDateLoader def loader_args(self, dates): (bound_expr,) = super( @@ -437,7 +371,8 @@ class BlazeDividendsByExDateLoaderNotInteractiveTestCase( return swap_resources_into_scope(bound_expr, {}) -class DividendsByPayDateTestCase(TestCase, EventLoaderCommonMixin): +class DividendsByPayDateTestCase(WithAssetFinder, ZiplineTestCase, + EventLoaderCommonMixin): """ Tests for loading the dividends by pay date data. """ @@ -453,60 +388,36 @@ class DividendsByPayDateTestCase(TestCase, EventLoaderCommonMixin): return range(0, 5) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - equity_info = 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'), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop([ANNOUNCEMENT_FIELD_NAME, - EX_DATE_FIELD_NAME], axis=1) - for sid, frame - in enumerate(dividends_cases)} - cls.finder = stack.enter_context( - tmp_asset_finder(equities=equity_info), - ) + def get_dataset(cls): + return {sid: + frame.drop([ANNOUNCEMENT_FIELD_NAME, + EX_DATE_FIELD_NAME], axis=1) + for sid, frame + in enumerate(dividends_cases)} - cls.loader_type = DividendsByPayDateLoader - - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = DividendsByPayDateLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) - - self.cols[NEXT_PAY_DATE] = get_vals_for_dates( - zip_with_dates_for_dates, num_days_between_for_dates, dates, - next_date_intervals, next_ex_and_pay_dates + cols = {} + cols[NEXT_PAY_DATE] = get_vals_for_dates( + zip_with_dates, next_ex_and_pay_dates, next_date_intervals, dates ) - self.cols[PREVIOUS_PAY_DATE] = get_vals_for_dates( - zip_with_dates_for_dates, num_days_between_for_dates, dates, - prev_date_intervals, prev_ex_and_pay_dates + cols[PREVIOUS_PAY_DATE] = get_vals_for_dates( + zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals, dates ) - self.cols[NEXT_AMOUNT] = get_vals_for_dates( - zip_with_floats_dates, num_days_between_dates, - dates, next_date_intervals, next_amounts + cols[NEXT_AMOUNT] = get_vals_for_dates( + zip_with_floats, next_amounts, next_date_intervals, dates ) - self.cols[PREVIOUS_AMOUNT] = get_vals_for_dates( - zip_with_floats_dates, num_days_between_dates, - dates, prev_date_intervals, prev_amounts + cols[PREVIOUS_AMOUNT] = get_vals_for_dates( + zip_with_floats, prev_amounts, prev_date_intervals, dates ) + return cols class BlazeDividendsByPayDateLoaderTestCase(DividendsByPayDateTestCase): - @classmethod - def setUpClass(cls): - super(BlazeDividendsByPayDateLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeDividendsByPayDateLoader + loader_type = BlazeDividendsByPayDateLoader def loader_args(self, dates): _, mapping = super( @@ -528,11 +439,6 @@ class BlazeDividendsByPayDateLoaderNotInteractiveTestCase( BlazeDividendsByPayDateLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - @classmethod - def setUpClass(cls): - super(BlazeDividendsByPayDateLoaderNotInteractiveTestCase, - cls).setUpClass() - cls.loader_type = BlazeDividendsByPayDateLoader def loader_args(self, dates): (bound_expr,) = super( diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index c67bfa03..0922a0c9 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -1,12 +1,8 @@ """ Tests for the reference loader for EarningsCalendar. """ -from functools import partial -from unittest import TestCase - import blaze as bz from blaze.compute.core import swap_resources_into_scope -from contextlib2 import ExitStack import pandas as pd from six import iteritems from .base import EventLoaderCommonMixin @@ -26,14 +22,13 @@ from zipline.pipeline.factors.events import ( BusinessDaysUntilNextEarnings, ) from zipline.pipeline.loaders.earnings import EarningsCalendarLoader -from zipline.pipeline.loaders.blaze import ( - BlazeEarningsCalendarLoader, -) -from zipline.utils.test_utils import ( - tmp_asset_finder, +from zipline.pipeline.loaders.blaze import BlazeEarningsCalendarLoader +from zipline.pipeline.loaders.utils import ( + get_values_for_date_ranges, + zip_with_dates ) -from zipline.testing import tmp_asset_finder +from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase earnings_cases = [ # K1--K2--A1--A2. @@ -63,8 +58,61 @@ earnings_cases = [ ), ] +next_date_intervals = [ + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-15'], + ['2014-01-16', '2014-01-20'], + ['2014-01-21', None]], + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-09'], + ['2014-01-10', '2014-01-15'], + ['2014-01-16', '2014-01-20'], + ['2014-01-21', None]], + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-10'], + ['2014-01-11', '2014-01-14'], + ['2014-01-15', '2014-01-20'], + ['2014-01-21', None]], + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-10'], + ['2014-01-11', '2014-01-15'], + ['2014-01-16', None]] +] -class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): +next_dates = [ + ['NaT', '2014-01-15', '2014-01-20', 'NaT'], + ['NaT', '2014-01-20', '2014-01-15', '2014-01-20', 'NaT'], + ['NaT', '2014-01-10', 'NaT', '2014-01-20', 'NaT'], + ['NaT', '2014-01-10', '2014-01-15', 'NaT'], + ['NaT'] +] + +prev_date_intervals = [ + [[None, '2014-01-14'], + ['2014-01-15', '2014-01-19'], + ['2014-01-20', None]], + [[None, '2014-01-14'], + ['2014-01-15', '2014-01-19'], + ['2014-01-20', None]], + [[None, '2014-01-09'], + ['2014-01-10', '2014-01-19'], + ['2014-01-20', None]], + [[None, '2014-01-09'], + ['2014-01-10', '2014-01-14'], + ['2014-01-15', None]] +] + +prev_dates = [ + ['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-15'], + ['NaT'] +] + + +class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase, + EventLoaderCommonMixin): """ Tests for loading the earnings announcement data. """ @@ -80,107 +128,53 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): return range(5) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - cls.cols = {} - cls.dataset = {sid: df for sid, df in enumerate(earnings_cases)} - cls.finder = stack.enter_context( - tmp_asset_finder(equities=cls.get_equity_info()), - ) + def get_dataset(cls): + return {sid: df for sid, df in enumerate(earnings_cases)} - cls.loader_type = EarningsCalendarLoader + loader_type = EarningsCalendarLoader def get_expected_next_event_dates(self, dates): - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) return pd.DataFrame({ - 0: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-15'] * - num_days_between_for_dates('2014-01-05', '2014-01-15') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-16', '2014-01-20') + - ['NaT'] * - num_days_between_for_dates('2014-01-21', None) - ), - 1: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-05', '2014-01-09') + - ['2014-01-15'] * - num_days_between_for_dates('2014-01-10', '2014-01-15') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-16', '2014-01-20') + - ['NaT'] * - num_days_between_for_dates('2014-01-21', None) - ), - 2: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-10'] * - num_days_between_for_dates('2014-01-05', '2014-01-10') + - ['NaT'] * - num_days_between_for_dates('2014-01-11', '2014-01-14') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-15', '2014-01-20') + - ['NaT'] * - num_days_between_for_dates('2014-01-21', None) - ), - 3: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-10'] * - num_days_between_for_dates('2014-01-05', '2014-01-10') + - ['2014-01-15'] * - num_days_between_for_dates('2014-01-11', '2014-01-15') + - ['NaT'] * - num_days_between_for_dates('2014-01-16', None) - ), - 4: zip_with_dates_for_dates(['NaT'] * - len(dates)), + 0: get_values_for_date_ranges(zip_with_dates, + next_dates[0], + next_date_intervals[0], + dates), + 1: get_values_for_date_ranges(zip_with_dates, + next_dates[1], + next_date_intervals[1], + dates), + 2: get_values_for_date_ranges(zip_with_dates, + next_dates[2], + next_date_intervals[2], + dates), + 3: get_values_for_date_ranges(zip_with_dates, + next_dates[3], + next_date_intervals[3], + dates), + 4: zip_with_dates(dates, ['NaT'] * len(dates)), }, index=dates) def get_expected_previous_event_dates(self, dates): - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) return pd.DataFrame({ - 0: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-14') + - ['2014-01-15'] * num_days_between_for_dates('2014-01-15', - '2014-01-19') + - ['2014-01-20'] * num_days_between_for_dates('2014-01-20', - None), - ), - 1: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-14') + - ['2014-01-15'] * num_days_between_for_dates('2014-01-15', - '2014-01-19') + - ['2014-01-20'] * num_days_between_for_dates('2014-01-20', - None), - ), - 2: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-09') + - ['2014-01-10'] * num_days_between_for_dates('2014-01-10', - '2014-01-19') + - ['2014-01-20'] * num_days_between_for_dates('2014-01-20', - None), - ), - 3: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-09') + - ['2014-01-10'] * num_days_between_for_dates('2014-01-10', - '2014-01-14') + - ['2014-01-15'] * num_days_between_for_dates('2014-01-15', - None), - ), - 4: zip_with_dates_for_dates(['NaT'] * len(dates)), + 0: get_values_for_date_ranges(zip_with_dates, + prev_dates[0], + prev_date_intervals[0], + dates), + 1: get_values_for_date_ranges(zip_with_dates, + prev_dates[1], + prev_date_intervals[1], + dates), + 2: get_values_for_date_ranges(zip_with_dates, + prev_dates[2], + prev_date_intervals[2], + dates), + 3: get_values_for_date_ranges(zip_with_dates, + prev_dates[3], + prev_date_intervals[3], + dates), + 4: zip_with_dates(dates, ['NaT'] * len(dates)), }, index=dates) - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() - def setup(self, dates): _expected_next_announce = self.get_expected_next_event_dates(dates) @@ -194,17 +188,16 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): _expected_previous_busday_offsets = self._compute_busday_offsets( _expected_previous_announce ) - self.cols[PREVIOUS_ANNOUNCEMENT] = _expected_previous_announce - self.cols[NEXT_ANNOUNCEMENT] = _expected_next_announce - self.cols[DAYS_TO_NEXT] = _expected_next_busday_offsets - self.cols[DAYS_SINCE_PREV] = _expected_previous_busday_offsets + cols = {} + cols[PREVIOUS_ANNOUNCEMENT] = _expected_previous_announce + cols[NEXT_ANNOUNCEMENT] = _expected_next_announce + cols[DAYS_TO_NEXT] = _expected_next_busday_offsets + cols[DAYS_SINCE_PREV] = _expected_previous_busday_offsets + return cols class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase): - @classmethod - def setUpClass(cls): - super(BlazeEarningsCalendarLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeEarningsCalendarLoader + loader_type = BlazeEarningsCalendarLoader def loader_args(self, dates): _, mapping = super( @@ -225,11 +218,6 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase( BlazeEarningsCalendarLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - @classmethod - def setUpClass(cls): - super(BlazeEarningsCalendarLoaderNotInteractiveTestCase, - cls).setUpClass() - cls.loader_type = BlazeEarningsCalendarLoader def loader_args(self, dates): (bound_expr,) = super( diff --git a/zipline/pipeline/loaders/blaze/dividends.py b/zipline/pipeline/loaders/blaze/dividends.py index 9971511e..86a332b9 100644 --- a/zipline/pipeline/loaders/blaze/dividends.py +++ b/zipline/pipeline/loaders/blaze/dividends.py @@ -6,10 +6,16 @@ from zipline.pipeline.common import ( SID_FIELD_NAME, TS_FIELD_NAME, ) -from zipline.pipeline.data.dividends import DividendsByExDate, \ - DividendsByAnnouncementDate, DividendsByPayDate -from zipline.pipeline.loaders.dividends import DividendsByAnnouncementDateLoader, \ - DividendsByPayDateLoader, DividendsByExDateLoader +from zipline.pipeline.data.dividends import ( + DividendsByExDate, + DividendsByAnnouncementDate, + DividendsByPayDate +) +from zipline.pipeline.loaders.dividends import ( + DividendsByAnnouncementDateLoader, + DividendsByPayDateLoader, + DividendsByExDateLoader +) from .events import BlazeEventsLoader diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py index 494c4992..c057f197 100644 --- a/zipline/pipeline/loaders/utils.py +++ b/zipline/pipeline/loaders/utils.py @@ -4,6 +4,7 @@ import numpy as np import pandas as pd from six import iteritems from six.moves import zip +from zipline.testing import num_days_in_range from zipline.utils.numpy_utils import NaTns @@ -46,9 +47,8 @@ def next_event_frame(events_by_sid, equity: np.full_like(dates, NaTns) for equity in events_by_sid } value_cols = { - equity: np.full(len(dates), missing_value, dtype=field_dtype) for equity - in - events_by_sid + equity: np.full(len(dates), missing_value, dtype=field_dtype) + for equity in events_by_sid } raw_dates = dates.values @@ -68,9 +68,9 @@ def next_event_frame(events_by_sid, (raw_dates <= event_date) ) value_mask = (event_date <= data) | (data == NaTns) - data_indeces = np.where(date_mask & value_mask) - data[data_indeces] = event_date - value_cols[equity][data_indeces] = value + data_indices = np.where(date_mask & value_mask) + data[data_indices] = event_date + value_cols[equity][data_indices] = value return pd.DataFrame(index=dates, data=value_cols) @@ -274,3 +274,53 @@ def check_data_query_args(data_query_time, data_query_tz): data_query_tz, ), ) + + +def zip_with_floats(dates, flts): + return pd.Series(flts, index=dates, dtype='float') + + +def num_days_between(dates, start_date, end_date): + return num_days_in_range(dates, start_date, end_date) + + +def zip_with_dates(index_dates, dts): + return pd.Series(pd.to_datetime(dts), index=index_dates) + + +def get_values_for_date_ranges(zip_date_index_with_vals, + vals_for_date_intervals, + date_intervals, + date_index): + """ + Returns a Series of values indexed by date based on values for the given + date intervals. + + Parameters + ---------- + zip_date_index_with_vals : callable + A function that takes in a list of dates and a list of values and + returns a pd.Series with the values indexed by the dates. + vals_for_date_intervals : list + A list of values for each date interval in `date_intervals`. + date_intervals : list + A list of pairs of dates, where each pair represents a date interval + that corresponds to the value at the same index in + `vals_for_date_intervals`. + date_index : DatetimeIndex + The DatetimeIndex containing all dates for which values were requested. + + Returns + ------- + date_index_with_vals : pd.Series + A Series indexed by the given DatetimeIndex and with values assigned + to dates based on the given date intervals. + """ + # Fill in given values for given date ranges. + return zip_date_index_with_vals( + date_index, + np.repeat(vals_for_date_intervals, + [num_days_between(date_index, *date_interval) + for date_interval in + date_intervals]), + ) diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index c8b8f04c..e4caf7ae 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -5,7 +5,7 @@ from logbook import NullHandler import pandas as pd from six import with_metaclass -from .core import tmp_asset_finder +from .core import tmp_asset_finder, make_simple_equity_info from ..finance.trading import TradingEnvironment from ..utils import tradingcalendar, factory from ..utils.final import FinalMeta, final @@ -177,7 +177,14 @@ class WithAssetFinder(object): def _make_info(cls): return None - make_equities_info = _make_info + @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_futures_info = _make_info make_exchanges_info = _make_info make_root_symbols_info = _make_info