diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index c76167ae..6b7218c0 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -1,14 +1,15 @@ """ 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 from nose_parameterized import parameterized -import pandas as pd import numpy as np +import pandas as pd from pandas.util.testing import assert_series_equal from six import iteritems @@ -23,85 +24,90 @@ from zipline.pipeline.factors.events import ( from zipline.pipeline.loaders.buyback_auth import \ CashBuybackAuthorizationsLoader, ShareBuybackAuthorizationsLoader from zipline.pipeline.loaders.blaze import ( + BlazeCashBuybackAuthorizationsLoader, + BlazeShareBuybackAuthorizationsLoader, BUYBACK_ANNOUNCEMENT_FIELD_NAME, - CashBuybackAuthorizationsLoader, SHARE_COUNT_FIELD_NAME, SID_FIELD_NAME, - ShareBuybackAuthorizationsLoader, TS_FIELD_NAME, - VALUE_FIELD_NAME + CASH_FIELD_NAME ) from zipline.utils.numpy_utils import make_datetime64D, np_NaT from zipline.utils.test_utils import ( - make_simple_equity_info, - tmp_asset_finder, gen_calendars, + make_simple_equity_info, num_days_in_range, + tmp_asset_finder, ) sids = A, B, C, D, E = range(5) equity_info = make_simple_equity_info( - sids, - start_date=pd.Timestamp('2013-01-01', tz='UTC'), - end_date=pd.Timestamp('2015-01-01', tz='UTC'), - ) + sids, + start_date=pd.Timestamp('2013-01-01', tz='UTC'), + end_date=pd.Timestamp('2015-01-01', tz='UTC'), +) buyback_authorizations = { - # K1--K2--A1--A2--SC1--SC2--V1--V2. - A: pd.DataFrame({ - "timestamp": pd.to_datetime(['2014-01-05', '2014-01-10']), - BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-15', - '2014-01-20']), - SHARE_COUNT_FIELD_NAME: [1, 15], - VALUE_FIELD_NAME: [10, 20] - }), - # K1--K2--E2--E1. - B: pd.DataFrame({ - "timestamp": pd.to_datetime(['2014-01-05', '2014-01-10']), - BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([ - '2014-01-20', '2014-01-15']), - SHARE_COUNT_FIELD_NAME: [7, 13], VALUE_FIELD_NAME: [10, 22] - }), - # K1--E1--K2--E2. - C: pd.DataFrame({ - "timestamp": pd.to_datetime(['2014-01-05', '2014-01-15']), - BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([ - '2014-01-10', '2014-01-20']), - SHARE_COUNT_FIELD_NAME: [3, 1], - VALUE_FIELD_NAME: [4, 7] - }), - # K1 == K2. - D: pd.DataFrame({ - "timestamp": pd.to_datetime(['2014-01-05'] * 2), - BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([ - '2014-01-10', '2014-01-15']), - SHARE_COUNT_FIELD_NAME: [6, 23], - VALUE_FIELD_NAME: [1, 2] - }), - E: pd.DataFrame( - columns=["timestamp", - BUYBACK_ANNOUNCEMENT_FIELD_NAME, - SHARE_COUNT_FIELD_NAME, - VALUE_FIELD_NAME], - dtype='datetime64[ns]' - ), - } - -param_dates = 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', + # K1--K2--A1--A2--SC1--SC2--V1--V2. + A: pd.DataFrame({ + "timestamp": pd.to_datetime(['2014-01-05', '2014-01-10']), + BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-15', + '2014-01-20']), + SHARE_COUNT_FIELD_NAME: [1, 15], + CASH_FIELD_NAME: [10, 20] + }), + # K1--K2--E2--E1. + B: pd.DataFrame({ + "timestamp": pd.to_datetime(['2014-01-05', '2014-01-10']), + BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([ + '2014-01-20', '2014-01-15' ]), - ) + SHARE_COUNT_FIELD_NAME: [7, 13], CASH_FIELD_NAME: [10, 22] + }), + # K1--E1--K2--E2. + C: pd.DataFrame({ + "timestamp": pd.to_datetime(['2014-01-05', '2014-01-15']), + BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([ + '2014-01-10', '2014-01-20' + ]), + SHARE_COUNT_FIELD_NAME: [3, 1], + CASH_FIELD_NAME: [4, 7] + }), + # K1 == K2. + D: pd.DataFrame({ + "timestamp": pd.to_datetime(['2014-01-05'] * 2), + BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([ + '2014-01-10', '2014-01-15' + ]), + SHARE_COUNT_FIELD_NAME: [6, 23], + CASH_FIELD_NAME: [1, 2] + }), + E: pd.DataFrame( + columns=["timestamp", + BUYBACK_ANNOUNCEMENT_FIELD_NAME, + SHARE_COUNT_FIELD_NAME, + CASH_FIELD_NAME], + dtype='datetime64[ns]' + ), +} + +# Must be a list - can't use generator since this needs to be used more than +# once. +param_dates = list(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', + ]), +)) -def zip_with_floats(flts, dates): +def zip_with_floats(dates, flts): return pd.Series(flts, index=dates).astype('float') @@ -109,30 +115,15 @@ def num_days_between(dates, start_date, end_date): return num_days_in_range(dates, start_date, end_date) -def zip_with_dates(dts, dates): - return pd.Series(pd.to_datetime(dts), index=dates) +def zip_with_dates(index_dates, dts): + return pd.Series(pd.to_datetime(dts), index=index_dates) -class BuybackAuthLoaderTestCase(TestCase): +class BuybackAuthLoaderCommonTest: """ - Tests for loading the earnings announcement data. + Tests for loading the buyback authorization announcement data. """ - @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - - cls.finder = stack.enter_context( - tmp_asset_finder(equities=equity_info), - ) - cls.cols = {} - cls.buyback_authorizations = None - - - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() - def loader_args(self, dates): """Construct the base buyback authorizations object to pass to the loader. @@ -149,54 +140,65 @@ class BuybackAuthLoaderTestCase(TestCase): """ return dates, self.buyback_authorizations - def setup(self, dates): + def setup_engine(self, dates): """ - Make a PipelineEngine and expectation functions for the given dates - calendar. + 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) + + def setup_expected_cols(self, dates): + """ + Make expectation functions for the given dates calendar. This exists to make it easy to test our various cases with critical dates missing from the calendar. """ - + num_days_between_for_dates = partial(num_days_between, dates) + zip_with_dates_for_dates = partial(zip_with_dates, dates) _expected_previous_buyback_announcement = pd.DataFrame({ - A: zip_with_dates( - ['NaT'] * num_days_between(dates, None, '2014-01-14') + - ['2014-01-15'] * num_days_between(dates, '2014-01-15', '2014-01-19') + - ['2014-01-20'] * num_days_between(dates, '2014-01-20', None), - dates + A: 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), ), - B: zip_with_dates( - ['NaT'] * num_days_between(dates, None, '2014-01-14') + - ['2014-01-15'] * num_days_between(dates, '2014-01-15', '2014-01-19') + - ['2014-01-20'] * num_days_between(dates, '2014-01-20', None), - dates + B: 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), ), - C: zip_with_dates( - ['NaT'] * num_days_between(dates, None, '2014-01-09') + - ['2014-01-10'] * num_days_between(dates, '2014-01-10', '2014-01-19') + - ['2014-01-20'] * num_days_between(dates, '2014-01-20', None), - dates + C: 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), ), - D: zip_with_dates( - ['NaT'] * num_days_between(dates, None, '2014-01-09') + - ['2014-01-10'] * num_days_between(dates, '2014-01-10', '2014-01-14') + - ['2014-01-15'] * num_days_between(dates, '2014-01-15', None), - dates + D: 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), ), - E: zip_with_dates(['NaT'] * len(dates), dates), + E: zip_with_dates_for_dates(['NaT'] * len(dates)), }, index=dates) _expected_previous_busday_offsets = self._compute_busday_offsets( _expected_previous_buyback_announcement ) - self.cols['previous_buyback_announcement'] = _expected_previous_buyback_announcement + # Common cols for buyback authorization datasets are announcement + # date and days since previous. + self.cols[ + 'previous_buyback_announcement' + ] = _expected_previous_buyback_announcement self.cols['days_since_prev'] = _expected_previous_busday_offsets - loader = self.loader_type(*self.loader_args(dates)) - engine = SimplePipelineEngine(lambda _: loader, dates, self.finder) - return engine - @staticmethod def _compute_busday_offsets(announcement_dates): """ @@ -234,7 +236,8 @@ class BuybackAuthLoaderTestCase(TestCase): ) def _test_compute_buyback_auth(self, dates): - engine = self.setup(dates) + engine = self.setup_engine(dates) + self.setup_expected_cols(dates) pipe = Pipeline( columns=self.pipeline_columns @@ -253,152 +256,247 @@ class BuybackAuthLoaderTestCase(TestCase): sid) -class ShareBuybackAuthLoaderTestCase(BuybackAuthLoaderTestCase): - buyback_authorizations = {sid: df.drop(VALUE_FIELD_NAME, 1) +class CashBuybackAuthLoaderTestCase(TestCase, BuybackAuthLoaderCommonTest): + """ + Test for cash buyback authorizations dataset. + """ + buyback_authorizations = {sid: df.drop(SHARE_COUNT_FIELD_NAME, 1) for sid, df in iteritems(buyback_authorizations)} pipeline_columns = { - 'previous_buyback_share_count': - ShareBuybackAuthorizations.previous_share_count.latest, - 'previous_buyback_announcement': - ShareBuybackAuthorizations.previous_announcement_date.latest, - 'days_since_prev': - BusinessDaysSincePreviousShareBuybackAuth(), - } + 'previous_buyback_cash': + CashBuybackAuthorizations.previous_value.latest, + 'previous_buyback_announcement': + CashBuybackAuthorizations.previous_announcement_date.latest, + 'days_since_prev': + BusinessDaysSincePreviousCashBuybackAuth(), + } @classmethod def setUpClass(cls): - super(ShareBuybackAuthLoaderTestCase, cls).setUpClass() + cls._cleanup_stack = stack = ExitStack() + cls.finder = stack.enter_context( + tmp_asset_finder(equities=equity_info), + ) + cls.cols = {} + cls.buyback_authorizations = buyback_authorizations + cls.loader_type = CashBuybackAuthorizationsLoader + + @classmethod + def tearDownClass(cls): + cls._cleanup_stack.close() + + def setup(self, dates): + zip_with_floats_dates = partial(zip_with_floats, dates) + num_days_between_dates = partial(num_days_between, dates) + super(CashBuybackAuthLoaderTestCase, self).setup_expected_cols(dates) + _expected_previous_cash = pd.DataFrame({ + # TODO if the next knowledge date is 10, why is the range + # until 15? + A: zip_with_floats_dates( + ['NaN'] * num_days_between(dates, None, '2014-01-14') + + [10] * num_days_between_dates('2014-01-15', '2014-01-19') + + [20] * num_days_between_dates('2014-01-20', None) + ), + B: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-14') + + [22] * num_days_between_dates('2014-01-15', '2014-01-19') + + [10] * num_days_between_dates('2014-01-20', None) + ), + C: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-09') + + [4] * num_days_between_dates('2014-01-10', '2014-01-19') + + [7] * num_days_between_dates('2014-01-20', None) + ), + D: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-09') + + [1] * num_days_between_dates('2014-01-10', '2014-01-14') + + [2] * num_days_between_dates('2014-01-15', None) + ), + E: zip_with_floats_dates(['NaN'] * len(dates)), + }, index=dates) + self.cols['previous_buyback_cash'] = _expected_previous_cash + + @parameterized.expand(param_dates) + def test_compute_cash_buyback_auth(self, dates): + self._test_compute_buyback_auth(dates) + + +class ShareBuybackAuthLoaderTestCase(BuybackAuthLoaderCommonTest, TestCase): + """ + Test for share buyback authorizations dataset. + """ + buyback_authorizations = {sid: df.drop(CASH_FIELD_NAME, 1) + for sid, df in iteritems(buyback_authorizations)} + pipeline_columns = { + 'previous_buyback_share_count': + ShareBuybackAuthorizations.previous_share_count.latest, + 'previous_buyback_announcement': + ShareBuybackAuthorizations.previous_announcement_date.latest, + 'days_since_prev': + BusinessDaysSincePreviousShareBuybackAuth(), + } + + @classmethod + def setUpClass(cls): + cls._cleanup_stack = stack = ExitStack() + cls.finder = stack.enter_context( + tmp_asset_finder(equities=equity_info), + ) + cls.cols = {} cls.buyback_authorizations = buyback_authorizations cls.loader_type = ShareBuybackAuthorizationsLoader + @classmethod + def tearDownClass(cls): + cls._cleanup_stack.close() + def setup(self, dates): - engine = super(ShareBuybackAuthLoaderTestCase, self).setup(dates) + zip_with_floats_dates = partial(zip_with_floats, dates) + num_days_between_dates = partial(num_days_between, dates) + super(ShareBuybackAuthLoaderTestCase, self).setup_expected_cols(dates) _expected_previous_buyback_share_count = pd.DataFrame({ - A: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-14') + - [1] * num_days_between(dates, '2014-01-15', '2014-01-19') + - [15] * num_days_between(dates, '2014-01-20', None), dates), - B: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-14') + - [13] * num_days_between(dates, '2014-01-15', '2014-01-19') + - [7] * num_days_between(dates, '2014-01-20', None), dates), - C: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-09') + - [3] * num_days_between(dates, '2014-01-10', '2014-01-19') + - [1] * num_days_between(dates, '2014-01-20', None), dates), - D: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-09') + - [6] * num_days_between(dates, '2014-01-10', '2014-01-14') + - [23] * num_days_between(dates, '2014-01-15', None), dates), - E: zip_with_floats(['NaN'] * len(dates), dates), - }, index=dates) - self.cols['previous_buyback_share_count'] = _expected_previous_buyback_share_count - return engine + A: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-14') + + [1] * num_days_between_dates('2014-01-15', '2014-01-19') + + [15] * num_days_between_dates('2014-01-20', None) + ), + B: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-14') + + [13] * num_days_between_dates('2014-01-15', '2014-01-19') + + [7] * num_days_between_dates('2014-01-20', None) + ), + C: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-09') + + [3] * num_days_between_dates('2014-01-10', '2014-01-19') + + [1] * num_days_between_dates('2014-01-20', None) + ), + D: zip_with_floats_dates( + ['NaN'] * num_days_between_dates(None, '2014-01-09') + + [6] * num_days_between_dates('2014-01-10', '2014-01-14') + + [23] * num_days_between_dates('2014-01-15', None) + ), + E: zip_with_floats_dates(['NaN'] * len(dates)), + }, index=dates) + self.cols[ + 'previous_buyback_share_count' + ] = _expected_previous_buyback_share_count @parameterized.expand(param_dates) - def test_compute_buyback_auth(self, dates): + def test_compute_share_buyback_auth(self, dates): self._test_compute_buyback_auth(dates) -class CashBuybackAuthLoaderTestCase(BuybackAuthLoaderTestCase): - buyback_authorizations = {sid: df.drop(SHARE_COUNT_FIELD_NAME, 1) - for sid, df in iteritems(buyback_authorizations)} - pipeline_columns = { - 'previous_buyback_value': - CashBuybackAuthorizations.previous_value.latest, - 'previous_buyback_announcement': - CashBuybackAuthorizations.previous_announcement_date.latest, - 'days_since_prev': - BusinessDaysSincePreviousCashBuybackAuth(), - } +def mapping_to_df(mapping): + return (bz.Data(pd.concat( + pd.DataFrame({ + BUYBACK_ANNOUNCEMENT_FIELD_NAME: + frame[BUYBACK_ANNOUNCEMENT_FIELD_NAME], + SHARE_COUNT_FIELD_NAME: + frame[SHARE_COUNT_FIELD_NAME], + CASH_FIELD_NAME: + frame[CASH_FIELD_NAME], + TS_FIELD_NAME: + frame[TS_FIELD_NAME], + SID_FIELD_NAME: sid, + }) + for sid, frame in iteritems(mapping) + ).reset_index(drop=True)),) + +class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): + """ Test case for loading via blaze. + """ @classmethod def setUpClass(cls): - super(CashBuybackAuthLoaderTestCase, cls).setUpClass() - cls.buyback_authorizations = buyback_authorizations - cls.loader_type = CashBuybackAuthLoaderTestCase + super(BlazeCashBuybackAuthLoaderTestCase, cls).setUpClass() + cls.loader_type = BlazeCashBuybackAuthorizationsLoader - def setup(self, dates): - engine = super(ShareBuybackAuthLoaderTestCase, self).setup(dates) - _expected_previous_value = pd.DataFrame({ - # TODO if the next knowledge date is 10, why is the range - # until 15? - A: zip_with_floats( - ['NaN'] * num_days_between(dates, None, '2014-01-14') + - [10] * num_days_between(dates, '2014-01-15', '2014-01-19') + - [20] * num_days_between(dates, '2014-01-20', None), dates), - B: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-14') + - [22] * num_days_between(dates, '2014-01-15', '2014-01-19') + - [10] * num_days_between(dates, '2014-01-20', None), dates), - C: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-09') + - [4] * num_days_between(dates, '2014-01-10', '2014-01-19') + - [7] * num_days_between(dates, '2014-01-20', None), dates), - D: zip_with_floats(['NaN'] * num_days_between(dates, None, '2014-01-09') + - [1] * num_days_between(dates, '2014-01-10', '2014-01-14') + - [2] * num_days_between(dates, '2014-01-15', None), dates), - E: zip_with_floats(['NaN'] * len(dates), dates), - }, index=dates) - self.cols['previous_buyback_value'] = _expected_previous_value - return engine - - @parameterized.expand(param_dates) - def test_compute_buyback_auth(self, dates): - self._test_compute_buyback_auth(dates) + def loader_args(self, dates): + _, mapping = super( + BlazeCashBuybackAuthLoaderTestCase, + self, + ).loader_args(dates) + return mapping_to_df(mapping) -# class BlazeBuybackAuthLoaderTestCase(BuybackAuthLoaderTestCase): -# loader_type = BlazeBuybackAuthorizationsLoader -# -# def loader_args(self, dates): -# _, mapping = super( -# BlazeBuybackAuthLoaderTestCase, -# self, -# ).loader_args(dates) -# return (bz.Data(pd.concat( -# pd.DataFrame({ -# BUYBACK_ANNOUNCEMENT_FIELD_NAME: -# frame[BUYBACK_ANNOUNCEMENT_FIELD_NAME], -# SHARE_COUNT_FIELD_NAME: frame[SHARE_COUNT_FIELD_NAME], -# VALUE_FIELD_NAME: frame[VALUE_FIELD_NAME], -# TS_FIELD_NAME: frame.index, -# SID_FIELD_NAME: sid, -# }) -# for sid, frame in iteritems(mapping) -# ).reset_index(drop=True)),) -# -# -# class BlazeEarningsCalendarLoaderNotInteractiveTestCase( -# BlazeBuybackAuthLoaderTestCase): -# """Test case for passing a non-interactive symbol and a dict of resources. -# """ -# def loader_args(self, dates): -# (bound_expr,) = super( -# BlazeEarningsCalendarLoaderNotInteractiveTestCase, -# self, -# ).loader_args(dates) -# return swap_resources_into_scope(bound_expr, {}) -# -# -# class BuybackAuthLoaderInferTimestampTestCase(TestCase): -# def test_infer_timestamp(self): -# dtx = pd.date_range('2014-01-01', '2014-01-10') -# events_by_sid = { -# 0: pd.DataFrame({BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx}), -# 1: pd.DataFrame( -# {BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.Series(dtx, dtx)}, -# index=dtx -# ) -# } -# loader = BuybackAuthorizationsLoader( -# dtx, -# events_by_sid, -# infer_timestamps=True, -# ) -# self.assertEqual( -# loader.events_by_sid.keys(), -# events_by_sid.keys(), -# ) -# assert_series_equal( -# loader.events_by_sid[0][BUYBACK_ANNOUNCEMENT_FIELD_NAME], -# pd.Series(index=[dtx[0]] * 10, data=dtx), -# ) -# assert_series_equal( -# loader.events_by_sid[1][BUYBACK_ANNOUNCEMENT_FIELD_NAME], -# events_by_sid[1][BUYBACK_ANNOUNCEMENT_FIELD_NAME], -# ) +class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase): + """ Test case for loading via blaze. + """ + @classmethod + def setUpClass(cls): + super(BlazeShareBuybackAuthLoaderTestCase, cls).setUpClass() + cls.loader_type = BlazeShareBuybackAuthorizationsLoader + + def loader_args(self, dates): + _, mapping = super( + BlazeShareBuybackAuthLoaderTestCase, + self, + ).loader_args(dates) + return mapping_to_df(mapping) + + +class BlazeShareBuybackAuthLoaderNotInteractiveTestCase( + BlazeShareBuybackAuthLoaderTestCase): + """Test case for passing a non-interactive symbol and a dict of resources. + """ + def loader_args(self, dates): + (bound_expr,) = super( + BlazeShareBuybackAuthLoaderNotInteractiveTestCase, + self, + ).loader_args(dates) + return swap_resources_into_scope(bound_expr, {}) + + +class BlazeCashBuybackAuthLoaderNotInteractiveTestCase( + BlazeCashBuybackAuthLoaderTestCase): + """Test case for passing a non-interactive symbol and a dict of resources. + """ + def loader_args(self, dates): + (bound_expr,) = super( + BlazeCashBuybackAuthLoaderNotInteractiveTestCase, + self, + ).loader_args(dates) + return swap_resources_into_scope(bound_expr, {}) + + +class BuybackAuthLoaderInferTimestampTestCase(TestCase): + @parameterized.expand([[CashBuybackAuthorizationsLoader], + [ShareBuybackAuthorizationsLoader]]) + def test_infer_timestamp(self, loader): + dtx = pd.date_range('2014-01-01', '2014-01-10') + events_by_sid = { + # No timestamp column - should index by first given date + 0: pd.DataFrame({BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx}), + # timestamp column exists - should index by it + 1: pd.DataFrame( + {BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx, + TS_FIELD_NAME: dtx} + ) + } + loader = loader( + dtx, + events_by_sid, + infer_timestamps=True, + ) + self.assertEqual( + loader.events_by_sid.keys(), + events_by_sid.keys(), + ) + + # Check that index by first given date has been added + assert_series_equal( + loader.events_by_sid[0][BUYBACK_ANNOUNCEMENT_FIELD_NAME], + pd.Series(index=[dtx[0]] * 10, + data=dtx, + name=BUYBACK_ANNOUNCEMENT_FIELD_NAME), + ) + + # Check that timestamp column was turned into index + modified_events_by_sid_date_col = pd.Series(data=np.array( + events_by_sid[1][BUYBACK_ANNOUNCEMENT_FIELD_NAME]), + index=events_by_sid[1][TS_FIELD_NAME], + name=BUYBACK_ANNOUNCEMENT_FIELD_NAME) + assert_series_equal( + loader.events_by_sid[1][BUYBACK_ANNOUNCEMENT_FIELD_NAME], + modified_events_by_sid_date_col, + ) diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index 4f78fb68..aecc2c88 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -7,8 +7,8 @@ import blaze as bz from blaze.compute.core import swap_resources_into_scope from contextlib2 import ExitStack from nose_parameterized import parameterized -import pandas as pd import numpy as np +import pandas as pd from pandas.util.testing import assert_series_equal from six import iteritems @@ -16,8 +16,8 @@ from zipline.pipeline import Pipeline from zipline.pipeline.data import EarningsCalendar from zipline.pipeline.engine import SimplePipelineEngine from zipline.pipeline.factors.events import ( - BusinessDaysUntilNextEarnings, BusinessDaysSincePreviousEarnings, + BusinessDaysUntilNextEarnings, ) from zipline.pipeline.loaders.earnings import EarningsCalendarLoader from zipline.pipeline.loaders.blaze import ( @@ -28,11 +28,10 @@ from zipline.pipeline.loaders.blaze import ( ) from zipline.utils.numpy_utils import make_datetime64D, NaTD from zipline.utils.test_utils import ( - make_simple_equity_info, - tmp_asset_finder, gen_calendars, - to_series, + make_simple_equity_info, num_days_in_range, + tmp_asset_finder, ) @@ -121,8 +120,7 @@ class EarningsCalendarLoaderTestCase(TestCase): def zip_with_dates(dts): return pd.Series(pd.to_datetime(dts), index=dates) - # TODO: tests will break because I now need mappings of sid -> - # dataframe instead of sid -> series + _expected_next_announce = pd.DataFrame({ A: zip_with_dates( ['NaT'] * num_days_between(None, '2014-01-04') + @@ -374,7 +372,9 @@ class EarningsCalendarLoaderInferTimestampTestCase(TestCase): dtx = pd.date_range('2014-01-01', '2014-01-10') announcement_dates = { 0: pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}), - 1: pd.DataFrame({TS_FIELD_NAME: dtx, ANNOUNCEMENT_FIELD_NAME: dtx}), + 1: pd.DataFrame( + {TS_FIELD_NAME: dtx, ANNOUNCEMENT_FIELD_NAME: dtx} + ), } loader = EarningsCalendarLoader( dtx, @@ -387,13 +387,15 @@ class EarningsCalendarLoaderInferTimestampTestCase(TestCase): ) assert_series_equal( pd.Series(loader.events_by_sid[0][ANNOUNCEMENT_FIELD_NAME]), - pd.Series(index=[dtx[0]] * 10, data=dtx, + pd.Series(index=[dtx[0]] * 10, + data=dtx, name=ANNOUNCEMENT_FIELD_NAME), ) assert_series_equal( pd.Series(loader.events_by_sid[1][ANNOUNCEMENT_FIELD_NAME]), pd.Series(index=announcement_dates[1][TS_FIELD_NAME], - data=np.array(announcement_dates[1][ - ANNOUNCEMENT_FIELD_NAME]), + data=np.array( + announcement_dates[1][ANNOUNCEMENT_FIELD_NAME] + ), name=ANNOUNCEMENT_FIELD_NAME) ) diff --git a/zipline/pipeline/data/buyback_auth.py b/zipline/pipeline/data/buyback_auth.py index 7c1cf952..8541d2a8 100644 --- a/zipline/pipeline/data/buyback_auth.py +++ b/zipline/pipeline/data/buyback_auth.py @@ -1,5 +1,5 @@ """ -Dataset representing dates of upcoming earnings. +Datasets representing dates of recently announced buyback authorizations. """ from zipline.utils.numpy_utils import datetime64ns_dtype, float64_dtype @@ -8,12 +8,17 @@ from .dataset import Column, DataSet class CashBuybackAuthorizations(DataSet): """ - Dataset representing dates of recently announced buyback authorization. + Dataset representing dates of recently announced cash buyback + authorizations. """ previous_value = Column(float64_dtype) previous_announcement_date = Column(datetime64ns_dtype) class ShareBuybackAuthorizations(DataSet): + """ + Dataset representing dates of recently announced share buyback + authorizations. + """ previous_share_count = Column(float64_dtype) previous_announcement_date = Column(datetime64ns_dtype) diff --git a/zipline/pipeline/factors/events.py b/zipline/pipeline/factors/events.py index 481e0e8a..127aae29 100644 --- a/zipline/pipeline/factors/events.py +++ b/zipline/pipeline/factors/events.py @@ -27,9 +27,9 @@ class BusinessDaysSincePreviousEvents(Factor): This doesn't use trading days for symmetry with BusinessDaysUntilNextEarnings. - Assets which announced or will announce the event today will produce a value - of 0.0. Assets that announced the event on the previous business day will - produce a value of 1.0. + Assets which announced or will announce the event today will produce a + value of 0.0. Assets that announced the event on the previous business + day will produce a value of 1.0. Assets for which the event date is `NaT` will produce a value of `NaN`. """ @@ -108,14 +108,16 @@ class BusinessDaysSincePreviousEarnings(BusinessDaysSincePreviousEvents): inputs = [EarningsCalendar.previous_announcement] -class BusinessDaysSincePreviousCashBuybackAuth(BusinessDaysSincePreviousEvents): +class BusinessDaysSincePreviousCashBuybackAuth( + BusinessDaysSincePreviousEvents +): """ Factor returning the number of **business days** (not trading days!) since the most recent cash buyback authorization for each asset. See Also -------- - zipline.pipeline.factors.BusinessDaysUntilNextEarnings + zipline.pipeline.factors.BusinessDaysSincePreviousCashBuybackAuth """ inputs = [CashBuybackAuthorizations.previous_announcement_date] @@ -130,6 +132,6 @@ class BusinessDaysSincePreviousShareBuybackAuth( See Also -------- - zipline.pipeline.factors.BusinessDaysUntilNextEarnings + zipline.pipeline.factors.BusinessDaysSincePreviousShareBuybackAuth """ inputs = [ShareBuybackAuthorizations.previous_announcement_date] diff --git a/zipline/pipeline/loaders/blaze/__init__.py b/zipline/pipeline/loaders/blaze/__init__.py index 9702d5ea..301cbc7d 100644 --- a/zipline/pipeline/loaders/blaze/__init__.py +++ b/zipline/pipeline/loaders/blaze/__init__.py @@ -1,6 +1,7 @@ + from .buyback_auth import ( - CashBuybackAuthorizationsLoader, - ShareBuybackAuthorizationsLoader + BlazeCashBuybackAuthorizationsLoader, + BlazeShareBuybackAuthorizationsLoader ) from .core import ( AD_FIELD_NAME, @@ -14,7 +15,7 @@ from .core import ( from .buyback_auth import ( BUYBACK_ANNOUNCEMENT_FIELD_NAME, SHARE_COUNT_FIELD_NAME, - VALUE_FIELD_NAME + CASH_FIELD_NAME ) from .earnings import ( ANNOUNCEMENT_FIELD_NAME, @@ -24,16 +25,16 @@ from .earnings import ( __all__ = ( 'AD_FIELD_NAME', 'ANNOUNCEMENT_FIELD_NAME', + 'BlazeCashBuybackAuthorizationsLoader', 'BlazeEarningsCalendarLoader', 'BlazeLoader', + 'BlazeShareBuybackAuthorizationsLoader', 'BUYBACK_ANNOUNCEMENT_FIELD_NAME', - 'CashBuybackAuthorizationsLoader', 'NoDeltasWarning', 'SHARE_COUNT_FIELD_NAME', 'SID_FIELD_NAME', - 'ShareBuybackAuthorizationsLoader', 'TS_FIELD_NAME', - 'VALUE_FIELD_NAME', + 'CASH_FIELD_NAME', 'from_blaze', 'global_loader', ) diff --git a/zipline/pipeline/loaders/blaze/buyback_auth.py b/zipline/pipeline/loaders/blaze/buyback_auth.py index b98a03aa..e92c3b9c 100644 --- a/zipline/pipeline/loaders/blaze/buyback_auth.py +++ b/zipline/pipeline/loaders/blaze/buyback_auth.py @@ -5,19 +5,17 @@ from .core import ( from zipline.pipeline.data import (CashBuybackAuthorizations, ShareBuybackAuthorizations) from zipline.pipeline.loaders.buyback_auth import ( + BUYBACK_ANNOUNCEMENT_FIELD_NAME, CashBuybackAuthorizationsLoader, - ShareBuybackAuthorizationsLoader + CASH_FIELD_NAME, + ShareBuybackAuthorizationsLoader, + SHARE_COUNT_FIELD_NAME ) from .events import BlazeEventsCalendarLoader -BUYBACK_ANNOUNCEMENT_FIELD_NAME = 'buyback_dates' -SHARE_COUNT_FIELD_NAME = 'share_counts' -VALUE_FIELD_NAME = 'values' - - class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): - """A pipeline loader for the ``BuybackAuth`` dataset that loads + """A pipeline loader for the ``CashBuybackAuthorizations`` dataset that loads data from a blaze expression. Parameters @@ -32,6 +30,10 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): The time to use for the data query cutoff. data_query_tz : tzinfo or str The timezeone to use for the data query cutoff. + dataset: DataSet + The DataSet object for which this loader loads data. + loader: EventsLoader + The reference loader to use for this dataset. Notes ----- @@ -41,12 +43,12 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): {SID_FIELD_NAME}: int64, {TS_FIELD_NAME}: datetime, {BUYBACK_ANNOUNCEMENT_FIELD_NAME}: ?datetime, - {VALUE_FIELD_NAME}: ?float64 + {CASH_FIELD_NAME}: ?float64 }} Where each row of the table is a record including the sid to identify the company, the timestamp where we learned about the announcement, the - date when the buyback was announced, the share count, and the value. + date when the buyback was announced, the share count, and the cash amount. If the '{TS_FIELD_NAME}' field is not included it is assumed that we start the backtest with knowledge of all announcements. @@ -55,28 +57,39 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): TS_FIELD_NAME=TS_FIELD_NAME, SID_FIELD_NAME=SID_FIELD_NAME, BUYBACK_ANNOUNCEMENT_FIELD_NAME=BUYBACK_ANNOUNCEMENT_FIELD_NAME, - VALUE_FIELD_NAME=VALUE_FIELD_NAME + CASH_FIELD_NAME=CASH_FIELD_NAME ) _expected_fields = frozenset({ TS_FIELD_NAME, SID_FIELD_NAME, BUYBACK_ANNOUNCEMENT_FIELD_NAME, - VALUE_FIELD_NAME + CASH_FIELD_NAME }) def __init__(self, expr, + resources=None, + odo_kwargs=None, + data_query_time=None, + data_query_tz=None, dataset=CashBuybackAuthorizations, loader=CashBuybackAuthorizationsLoader, **kwargs): super( BlazeCashBuybackAuthorizationsLoader, self - ).__init__(expr, dataset=dataset, loader=loader, **kwargs) + ).__init__(expr, + resources=resources, + odo_kwargs=odo_kwargs, + data_query_time=data_query_time, + data_query_tz=data_query_tz, + dataset=dataset, + loader=loader, + **kwargs) class BlazeShareBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): - """A pipeline loader for the ``BuybackAuth`` dataset that loads + """A pipeline loader for the ``ShareBuybackAuthorizations`` dataset that loads data from a blaze expression. Parameters @@ -91,6 +104,10 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): The time to use for the data query cutoff. data_query_tz : tzinfo or str The timezeone to use for the data query cutoff. + dataset: DataSet + The DataSet object for which this loader loads data. + loader: EventsLoader + The reference loader to use for this dataset. Notes ----- @@ -126,9 +143,20 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsCalendarLoader): def __init__(self, expr, + resources=None, + odo_kwargs=None, + data_query_time=None, + data_query_tz=None, dataset=ShareBuybackAuthorizations, loader=ShareBuybackAuthorizationsLoader, **kwargs): super( BlazeShareBuybackAuthorizationsLoader, self - ).__init__(expr, dataset=dataset, loader=loader, **kwargs) + ).__init__(expr, + resources=resources, + odo_kwargs=odo_kwargs, + data_query_time=data_query_time, + data_query_tz=data_query_tz, + dataset=dataset, + loader=loader, + **kwargs) diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 287949da..a08c2fd8 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -24,6 +24,10 @@ class BlazeEarningsCalendarLoader(BlazeEventsCalendarLoader): The time to use for the data query cutoff. data_query_tz : tzinfo or str The timezeone to use for the data query cutoff. + dataset: DataSet + The DataSet object for which this loader loads data. + loader: EventsLoader + The reference loader to use for this dataset. Notes ----- diff --git a/zipline/pipeline/loaders/blaze/events.py b/zipline/pipeline/loaders/blaze/events.py index c74cad5b..8196377c 100644 --- a/zipline/pipeline/loaders/blaze/events.py +++ b/zipline/pipeline/loaders/blaze/events.py @@ -16,7 +16,6 @@ from zipline.utils.input_validation import ensure_timezone, optionally from zipline.utils.preprocess import preprocess - class BlazeEventsCalendarLoader(PipelineLoader): """An abstract pipeline loader for the events datasets that loads data from a blaze expression. @@ -33,7 +32,10 @@ class BlazeEventsCalendarLoader(PipelineLoader): The time to use for the data query cutoff. data_query_tz : tzinfo or str The timezeone to use for the data query cutoff. - + dataset : DataSet + The DataSet object for which this loader loads data. + concrete_loader : + The concrete loader to use for loading data into specified columns. Notes ----- The expression should have a tabular dshape of:: @@ -60,7 +62,7 @@ class BlazeEventsCalendarLoader(PipelineLoader): data_query_time=None, data_query_tz=None, dataset=None, - loader=None): + concrete_loader=None): dshape = expr.dshape if not istabular(dshape): @@ -78,7 +80,7 @@ class BlazeEventsCalendarLoader(PipelineLoader): check_data_query_args(data_query_time, data_query_tz) self._data_query_time = data_query_time self._data_query_tz = data_query_tz - self._loader = loader + self._concrete_loader = concrete_loader def load_adjusted_array(self, columns, dates, assets, mask): data_query_time = self._data_query_time @@ -110,7 +112,7 @@ class BlazeEventsCalendarLoader(PipelineLoader): ts_field=TS_FIELD_NAME, ) gb = raw.groupby(SID_FIELD_NAME) - return self._loader( + return self._concrete_loader( dates, self.prepare_data(raw, gb), dataset=self._dataset, diff --git a/zipline/pipeline/loaders/buyback_auth.py b/zipline/pipeline/loaders/buyback_auth.py index 02727e53..038a89c0 100644 --- a/zipline/pipeline/loaders/buyback_auth.py +++ b/zipline/pipeline/loaders/buyback_auth.py @@ -2,30 +2,26 @@ Reference implementation for EarningsCalendar loaders. """ -from ..data.buyback_auth import CashBuybackAuthorizations, \ +from ..data.buyback_auth import ( + CashBuybackAuthorizations, ShareBuybackAuthorizations +) from events import EventsLoader from zipline.utils.memoize import lazyval BUYBACK_ANNOUNCEMENT_FIELD_NAME = 'buyback_dates' SHARE_COUNT_FIELD_NAME = 'share_counts' -VALUE_FIELD_NAME = 'values' +CASH_FIELD_NAME = 'cash' -# TODO: split into 2 datasets - or just think about how to generalize since -# we will often have cases where we have a knowledge date and, optionally, -# a value for that event; having no value (like earnings) is a special case. class CashBuybackAuthorizationsLoader(EventsLoader): """ Reference loader for - :class:`zipline.pipeline.data.earnings.BuybackAuthorizations`. - - Does not currently support adjustments to the dates of known buyback - authorizations. + :class:`zipline.pipeline.data.earnings.CashBuybackAuthorizations`. events_by_sid: dict[sid -> pd.DataFrame(knowledge date, - event date, value)] + event date, cash value)] """ @@ -41,7 +37,6 @@ class CashBuybackAuthorizationsLoader(EventsLoader): dataset=dataset ) - def get_loader(self, column): """dispatch to the loader for ``column``. """ @@ -52,13 +47,12 @@ class CashBuybackAuthorizationsLoader(EventsLoader): else: raise ValueError("Don't know how to load column '%s'." % column) - @lazyval def previous_buyback_value_loader(self): return self._previous_event_value_loader( self.dataset.previous_value, BUYBACK_ANNOUNCEMENT_FIELD_NAME, - VALUE_FIELD_NAME + CASH_FIELD_NAME ) @lazyval @@ -72,13 +66,13 @@ class CashBuybackAuthorizationsLoader(EventsLoader): class ShareBuybackAuthorizationsLoader(EventsLoader): """ Reference loader for - :class:`zipline.pipeline.data.earnings.BuybackAuthorizations`. + :class:`zipline.pipeline.data.earnings.ShareBuybackAuthorizations`. Does not currently support adjustments to the dates of known buyback authorizations. events_by_sid: dict[sid -> pd.DataFrame(knowledge date, - event date, value)] + event date, share value)] """ @@ -94,7 +88,6 @@ class ShareBuybackAuthorizationsLoader(EventsLoader): dataset=dataset ) - def get_loader(self, column): """dispatch to the loader for ``column``. """ @@ -105,7 +98,6 @@ class ShareBuybackAuthorizationsLoader(EventsLoader): else: raise ValueError("Don't know how to load column '%s'." % column) - @lazyval def previous_buyback_share_count_loader(self): return self._previous_event_value_loader( diff --git a/zipline/pipeline/loaders/earnings.py b/zipline/pipeline/loaders/earnings.py index f93645d0..e2030430 100644 --- a/zipline/pipeline/loaders/earnings.py +++ b/zipline/pipeline/loaders/earnings.py @@ -2,8 +2,8 @@ Reference implementation for EarningsCalendar loaders. """ -from events import EventsLoader from ..data.earnings import EarningsCalendar +from events import EventsLoader from zipline.utils.memoize import lazyval ANNOUNCEMENT_FIELD_NAME = "announcement_date" diff --git a/zipline/pipeline/loaders/events.py b/zipline/pipeline/loaders/events.py index 8c6e2a5e..9920a729 100644 --- a/zipline/pipeline/loaders/events.py +++ b/zipline/pipeline/loaders/events.py @@ -1,4 +1,4 @@ -from abc import ABCMeta, abstractmethod +from abc import abstractmethod import numpy as np import pandas as pd @@ -22,22 +22,19 @@ class EventsLoader(PipelineLoader): ---------- all_dates : pd.DatetimeIndex Index of dates for which we can serve queries. - events_by_sid : dict[int -> pd.Series] - Dict mapping sids to objects representing dates on which events - occurred. + events_by_sid : dict[int -> pd.DataFrame] + Dict mapping sids to DataFrames representing dates on which events + occurred along with other associated values. - If a dict value is a Series, it's interpreted as a mapping from the - date on which we learned an announcement was coming to the date on - which the announcement was made. + If the DataFrames contain a "timestamp" column, that column is + interpreted as the date on which we learned about the event. - If a dict value is a DatetimeIndex, it's interpreted as just containing - the dates that announcements were made, and we assume we knew about the - announcement on all prior dates. This mode is only supported if - ``infer_timestamp`` is explicitly passed as a truthy value. + If the DataFrames do not contain a "timestamp" column, we assume we + knew about the event on all prior dates. This mode is only supported + if ``infer_timestamp`` is explicitly passed as a truthy value. infer_timestamps : bool, optional - Whether to allow passing ``DatetimeIndex`` values in - ``announcement_dates``. + Whether to allow omitting the "timestamp" column. """ def __init__(self, @@ -46,8 +43,9 @@ class EventsLoader(PipelineLoader): infer_timestamps=False, dataset=None): self.all_dates = all_dates - # TODO: why are we making a copy here? We end up with a copy that we - # modify and then don't use, and an unmodified original which we do use. + + # Do not modify the original in place, since it may be used for other + # purposes. self.events_by_sid = ( events_by_sid.copy() ) @@ -57,7 +55,8 @@ class EventsLoader(PipelineLoader): if "timestamp" not in v.columns: if not infer_timestamps: raise ValueError( - "Got DatetimeIndex of announcement dates for sid %d.\n" + "Got DataFrame without a 'timestamp' column for " + "sid %d.\n" "Pass `infer_timestamps=True` to use the first date in" " `all_dates` as implicit timestamp." ) @@ -68,11 +67,9 @@ class EventsLoader(PipelineLoader): self.dataset = dataset - @abstractmethod def get_loader(self): - raise NotImplementedError("EventsLoader must implement 'get_loader'.") - + raise NotImplementedError("Must implement 'get_loader'.") def load_adjusted_array(self, columns, dates, assets, mask): return merge( @@ -97,7 +94,9 @@ class EventsLoader(PipelineLoader): adjustments=None, ) - def _previous_event_date_loader(self, prev_date_field, event_date_field_name): + def _previous_event_date_loader(self, + prev_date_field, + event_date_field_name): return DataFrameLoader( prev_date_field, previous_date_frame( @@ -125,5 +124,3 @@ class EventsLoader(PipelineLoader): ), adjustments=None, ) - -