From 18d838be87ae68406db6f79f57c66f4921ef8502 Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Thu, 25 Feb 2016 13:52:16 -0500 Subject: [PATCH] MAINT: modify coersion logic. MAINT: make expected_cols class attributes. MAINT: make concrete_loader a class attribute. MAINT: drop sid column before creating loader. TST: add method to fill df with event cases. TST: move common case df into test_events.py. TST: improve tests and fix error message assertions. STY: fix whitespace. DOC: update docs. STY: fix style. MAINT: clean up STY: fix indentation. MAINT: use different assertion method to check error message for python3.4 compatibility. --- tests/pipeline/test_buyback_auth.py | 63 +++---- tests/pipeline/test_earnings.py | 45 +---- tests/pipeline/test_events.py | 156 +++++++++++++----- .../pipeline/loaders/blaze/buyback_auth.py | 24 +-- zipline/pipeline/loaders/blaze/earnings.py | 6 +- zipline/pipeline/loaders/blaze/events.py | 3 +- zipline/pipeline/loaders/buyback_auth.py | 13 +- zipline/pipeline/loaders/earnings.py | 6 +- zipline/pipeline/loaders/events.py | 83 ++++++---- 9 files changed, 212 insertions(+), 187 deletions(-) diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index 6d4ee10b..fbdd9797 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -7,10 +7,9 @@ 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 from six import iteritems -from tests.pipeline.test_events import param_dates, EventLoaderCommonTest +from tests.pipeline.test_events import EventLoaderCommonMixin, DATE_FIELD_NAME from zipline.pipeline.common import( BUYBACK_ANNOUNCEMENT_FIELD_NAME, @@ -42,54 +41,53 @@ from zipline.utils.test_utils import ( buyback_authorizations = [ # K1--K2--A1--A2. pd.DataFrame({ - TS_FIELD_NAME: 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--A2--A1. pd.DataFrame({ - TS_FIELD_NAME: 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] + SHARE_COUNT_FIELD_NAME: [7, 13], + CASH_FIELD_NAME: [10, 22] }), # K1--A1--K2--A2. pd.DataFrame({ - TS_FIELD_NAME: 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. pd.DataFrame({ - TS_FIELD_NAME: 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] }), pd.DataFrame( - columns=[TS_FIELD_NAME, - BUYBACK_ANNOUNCEMENT_FIELD_NAME, - SHARE_COUNT_FIELD_NAME, + columns=[SHARE_COUNT_FIELD_NAME, CASH_FIELD_NAME], dtype='datetime64[ns]' ), ] -class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest): +def create_buyback_auth_tst_frame(cases, field_to_drop): + buyback_auth_df = { + sid: + pd.concat([df, buyback_authorizations[sid]], axis=1).drop( + field_to_drop, 1) + for sid, df + in enumerate(case.rename(columns={DATE_FIELD_NAME: + BUYBACK_ANNOUNCEMENT_FIELD_NAME} + ) + for case in cases + ) + } + return buyback_auth_df + + +class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): """ Test for cash buyback authorizations dataset. """ pipeline_columns = { - (PREVIOUS_BUYBACK_CASH): + PREVIOUS_BUYBACK_CASH: CashBuybackAuthorizations.previous_value.latest, PREVIOUS_BUYBACK_ANNOUNCEMENT: CashBuybackAuthorizations.previous_announcement_date.latest, @@ -104,8 +102,8 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest): tmp_asset_finder(equities=cls.equity_info), ) cls.cols = {} - cls.dataset = {sid: df.drop(SHARE_COUNT_FIELD_NAME, 1) - for sid, df in enumerate(buyback_authorizations)} + cls.dataset = create_buyback_auth_tst_frame(cls.event_dates_cases, + SHARE_COUNT_FIELD_NAME) cls.loader_type = CashBuybackAuthorizationsLoader @classmethod @@ -145,12 +143,8 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest): self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] ) - @parameterized.expand(param_dates) - def test_compute_cash_buyback_auth(self, dates): - self._test_compute(dates) - -class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): +class ShareBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): """ Test for share buyback authorizations dataset. """ @@ -170,9 +164,8 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): tmp_asset_finder(equities=cls.equity_info), ) cls.cols = {} - cls.dataset = {sid: df.drop(CASH_FIELD_NAME, 1) - for sid, df in - enumerate(buyback_authorizations)} + cls.dataset = create_buyback_auth_tst_frame(cls.event_dates_cases, + CASH_FIELD_NAME) cls.loader_type = ShareBuybackAuthorizationsLoader @classmethod @@ -214,10 +207,6 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] ) - @parameterized.expand(param_dates) - def test_compute_share_buyback_auth(self, dates): - self._test_compute(dates) - class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): """ Test case for loading via blaze. diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index a4bce18b..ce39220d 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -6,10 +6,9 @@ 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 from six import iteritems -from tests.pipeline.test_events import EventLoaderCommonTest, param_dates +from tests.pipeline.test_events import EventLoaderCommonMixin, DATE_FIELD_NAME from zipline.pipeline.common import ( ANNOUNCEMENT_FIELD_NAME, @@ -34,39 +33,8 @@ from zipline.utils.test_utils import ( tmp_asset_finder, ) -earnings_dates = [ - # K1--K2--E1--E2. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-15', - '2014-01-20']) - }), - # K1--K2--E2--E1. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-20', - '2014-01-15']) - }), - # K1--E1--K2--E2. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-15']), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-10', - '2014-01-20']) - }), - # K1 == K2. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-10', - '2014-01-15']) - }), - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime([]), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([]) - }) -] - -class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonTest): +class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): """ Tests for loading the earnings announcement data. """ @@ -86,7 +54,10 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonTest): end_date=pd.Timestamp('2015-01-01', tz='UTC'), ) cls.cols = {} - cls.dataset = {sid: df for sid, df in enumerate(earnings_dates)} + cls.dataset = {sid: df for sid, df in enumerate( + case.rename( + columns={DATE_FIELD_NAME: ANNOUNCEMENT_FIELD_NAME} + ) for case in cls.event_dates_cases)} cls.finder = stack.enter_context( tmp_asset_finder(equities=equity_info), ) @@ -115,10 +86,6 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonTest): self.cols[DAYS_TO_NEXT] = _expected_next_busday_offsets self.cols[DAYS_SINCE_PREV] = _expected_previous_busday_offsets - @parameterized.expand(param_dates) - def test_compute_earnings(self, dates): - self._test_compute(dates) - class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase): @classmethod diff --git a/tests/pipeline/test_events.py b/tests/pipeline/test_events.py index 6c94b902..48cd4d07 100644 --- a/tests/pipeline/test_events.py +++ b/tests/pipeline/test_events.py @@ -3,11 +3,13 @@ Tests for setting up an EventsLoader and a BlazeEventsLoader. """ from functools import partial from nose_parameterized import parameterized +import re +from unittest import TestCase import blaze as bz import numpy as np import pandas as pd -from pandas.util.testing import assert_series_equal, TestCase +from pandas.util.testing import assert_series_equal from zipline.pipeline import SimplePipelineEngine, Pipeline from zipline.pipeline.common import ( @@ -18,21 +20,29 @@ from zipline.pipeline.common import ( from zipline.pipeline.data import DataSet, Column from zipline.pipeline.loaders.blaze.events import BlazeEventsLoader from zipline.pipeline.loaders.events import ( - BAD_DATA_FORMAT_ERROR, DF_NO_TS_NOT_INFER_TS_ERROR, DTINDEX_NOT_INFER_TS_ERROR, EventsLoader, SERIES_NO_DTINDEX_ERROR, WRONG_COLS_ERROR, + WRONG_MANY_COL_DATA_FORMAT_ERROR, + WRONG_SINGLE_COL_DATA_FORMAT_ERROR ) from zipline.utils.memoize import lazyval -from zipline.utils.numpy_utils import (datetime64ns_dtype, - NaTD, - make_datetime64D) -from zipline.utils.test_utils import gen_calendars, num_days_in_range, \ +from zipline.utils.numpy_utils import ( + datetime64ns_dtype, + NaTD, + make_datetime64D +) +from zipline.utils.test_utils import ( + gen_calendars, + num_days_in_range, make_simple_equity_info +) -ABSTRACT_METHODS_ERROR = 'abstract methods concrete_loader' +ABSTRACT_CONCRETE_LOADER_ERROR = 'abstract methods concrete_loader' +ABSTRACT_EXPECTED_COLS_ERROR = 'abstract methods expected_cols' +DATE_FIELD_NAME = "event_date" class EventDataSet(DataSet): @@ -40,6 +50,7 @@ class EventDataSet(DataSet): class EventDataSetLoader(EventsLoader): + expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME]) def __init__(self, all_dates, @@ -53,10 +64,6 @@ class EventDataSetLoader(EventsLoader): dataset=dataset, ) - @property - def expected_cols(self): - return frozenset([ANNOUNCEMENT_FIELD_NAME]) - @lazyval def previous_announcement_loader(self): return self._previous_event_date_loader( @@ -72,6 +79,12 @@ class EventDataSetLoader(EventsLoader): ) +# Test case just for catching an error when multiple columns are in the wrong +# data format, so no loader defined. +class EventDataSetLoaderMultipleExpectedCols(EventsLoader): + expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME, "other_field"]) + + class EventDataSetLoaderNoExpectedCols(EventsLoader): def __init__(self, @@ -90,32 +103,32 @@ class EventDataSetLoaderNoExpectedCols(EventsLoader): dtx = pd.date_range('2014-01-01', '2014-01-10') -def assert_loader_error(events_by_sid, error, msg, infer_timestamps): - with TestCase.assertRaises(error) as context: - EventDataSetLoader( - dtx, events_by_sid, infer_timestamps=infer_timestamps, - ) - TestCase.assertTrue(msg in context.exception) - - class EventLoaderTestCase(TestCase): + def assert_loader_error(self, events_by_sid, error, msg, + infer_timestamps, loader): + with self.assertRaisesRegexp(error, re.escape(msg)): + loader( + dtx, events_by_sid, infer_timestamps=infer_timestamps, + ) def test_no_expected_cols_defined(self): events_by_sid = {0: pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx})} - assert_loader_error(events_by_sid, TypeError, ABSTRACT_METHODS_ERROR, - True) + self.assert_loader_error(events_by_sid, TypeError, + ABSTRACT_EXPECTED_COLS_ERROR, + True, EventDataSetLoaderNoExpectedCols) def test_wrong_cols(self): wrong_col_name = 'some_other_col' # Test wrong cols (cols != expected) events_by_sid = {0: pd.DataFrame({wrong_col_name: dtx})} - assert_loader_error( + self.assert_loader_error( events_by_sid, ValueError, WRONG_COLS_ERROR.format( - expected_columns=EventDataSetLoader.expected_cols, + expected_columns=list(EventDataSetLoader.expected_cols), sid=0, - resulting_columns=wrong_col_name, + resulting_columns=[wrong_col_name], ), - True + True, + EventDataSetLoader ) @parameterized.expand([ @@ -125,9 +138,9 @@ class EventLoaderTestCase(TestCase): [pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx, TS_FIELD_NAME: dtx}), False], # DatetimeIndex with infer_timestamps = True - [pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME), True], + [pd.DatetimeIndex(dtx), True], # Series with DatetimeIndex as index and infer_timestamps = False - [pd.Series(dtx, index=dtx, name=ANNOUNCEMENT_FIELD_NAME), False] + [pd.Series(dtx, index=dtx), False] ]) def test_conversion_to_df(self, df, infer_timestamps): @@ -164,37 +177,61 @@ class EventLoaderTestCase(TestCase): DF_NO_TS_NOT_INFER_TS_ERROR.format( timestamp_column_name=TS_FIELD_NAME, sid=0 - ) + ), + EventDataSetLoader ], # DatetimeIndex with infer_timestamps = False [ pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME), False, - DTINDEX_NOT_INFER_TS_ERROR.format(sid=0) + DTINDEX_NOT_INFER_TS_ERROR.format(sid=0), + EventDataSetLoader ], # Series with DatetimeIndex as index and infer_timestamps = False [ pd.Series(dtx, name=ANNOUNCEMENT_FIELD_NAME), False, - SERIES_NO_DTINDEX_ERROR.format(sid=0) + SERIES_NO_DTINDEX_ERROR.format(sid=0), + EventDataSetLoader ], - # Some other data structure that is not expected + # Below, 2 cases repeated for infer_timestamps = True and False. + # Shouldn't make a difference in the outcome. + # We expected 1 column but got a data structure other than a + # DataFrame, Series, or DatetimeIndex [ - dtx, - False, - BAD_DATA_FORMAT_ERROR.format(sid=0) - ], - [ - dtx, + [dtx], True, - BAD_DATA_FORMAT_ERROR.format(sid=0) + WRONG_SINGLE_COL_DATA_FORMAT_ERROR.format(sid=0), + EventDataSetLoader + ], + # We expected multiple columns but got a data structure other + # than a DataFrame + [ + [dtx, dtx], + True, + WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=0), + EventDataSetLoaderMultipleExpectedCols + ], + [ + [dtx], + False, + WRONG_SINGLE_COL_DATA_FORMAT_ERROR.format(sid=0), + EventDataSetLoader + ], + # We expected multiple columns but got a data structure other + # than a DataFrame + [ + [dtx, dtx], + False, + WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=0), + EventDataSetLoaderMultipleExpectedCols ] ] ) - def test_bad_conversion_to_df(self, df, infer_timestamps, msg): + def test_bad_conversion_to_df(self, df, infer_timestamps, msg, loader): events_by_sid = {0: df} - assert_loader_error(events_by_sid, ValueError, msg, - infer_timestamps=infer_timestamps) + self.assert_loader_error(events_by_sid, ValueError, msg, + infer_timestamps, loader) class BlazeEventDataSetLoaderNoConcreteLoader(BlazeEventsLoader): @@ -212,7 +249,9 @@ class BlazeEventDataSetLoaderNoConcreteLoader(BlazeEventsLoader): class BlazeEventLoaderTestCase(TestCase): # Blaze loader: need to test failure if no concrete loader def test_no_concrete_loader_defined(self): - with TestCase.assertRaises(TypeError) as context: + with self.assertRaisesRegexp( + TypeError, re.escape(ABSTRACT_CONCRETE_LOADER_ERROR) + ): BlazeEventDataSetLoaderNoConcreteLoader( bz.Data( pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx, @@ -220,7 +259,6 @@ class BlazeEventLoaderTestCase(TestCase): }) ) ) - TestCase.assertTrue(ABSTRACT_METHODS_ERROR in context.exception) # Must be a list - can't use generator since this needs to be used more than @@ -237,7 +275,7 @@ param_dates = list(gen_calendars( )) -class EventLoaderCommonTest(object): +class EventLoaderCommonMixin(object): sids = A, B, C, D, E = range(5) equity_info = make_simple_equity_info( sids, @@ -245,6 +283,33 @@ class EventLoaderCommonTest(object): end_date=pd.Timestamp('2015-01-01', tz='UTC'), ) + event_dates_cases = [ + # K1--K2--E1--E2. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), + DATE_FIELD_NAME: pd.to_datetime(['2014-01-15', '2014-01-20']) + }), + # K1--K2--E2--E1. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), + DATE_FIELD_NAME: pd.to_datetime(['2014-01-20', '2014-01-15']) + }), + # K1--E1--K2--E2. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-15']), + DATE_FIELD_NAME: pd.to_datetime(['2014-01-10', '2014-01-20']) + }), + # K1 == K2. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2), + DATE_FIELD_NAME: pd.to_datetime(['2014-01-10', '2014-01-15']) + }), + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime([]), + DATE_FIELD_NAME: pd.to_datetime([]) + }) + ] + def zip_with_floats(self, dates, flts): return pd.Series(flts, index=dates).astype('float') @@ -399,7 +464,8 @@ class EventLoaderCommonTest(object): index=announcement_dates.index, ) - def _test_compute(self, dates): + @parameterized.expand(param_dates) + def test_compute(self, dates): engine = self.setup_engine(dates) self.setup(dates) diff --git a/zipline/pipeline/loaders/blaze/buyback_auth.py b/zipline/pipeline/loaders/blaze/buyback_auth.py index 52305e30..cd83e1d1 100644 --- a/zipline/pipeline/loaders/blaze/buyback_auth.py +++ b/zipline/pipeline/loaders/blaze/buyback_auth.py @@ -17,8 +17,8 @@ from .events import BlazeEventsLoader class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader): - """A pipeline loader for the ``CashBuybackAuthorizations`` dataset that loads - data from a blaze expression. + """A pipeline loader for the ``CashBuybackAuthorizations`` dataset that + loads data from a blaze expression. Parameters ---------- @@ -34,8 +34,6 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader): The timezeone to use for the data query cutoff. dataset: DataSet The DataSet object for which this loader loads data. - concrete_loader: EventsLoader - The reference loader to use for this dataset. Notes ----- @@ -69,6 +67,8 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader): CASH_FIELD_NAME }) + concrete_loader = CashBuybackAuthorizationsLoader + def __init__(self, expr, resources=None, @@ -87,14 +87,10 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader): dataset=dataset, **kwargs) - @property - def concrete_loader(self): - return CashBuybackAuthorizationsLoader - class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): - """A pipeline loader for the ``ShareBuybackAuthorizations`` dataset that loads - data from a blaze expression. + """A pipeline loader for the ``ShareBuybackAuthorizations`` dataset that + loads data from a blaze expression. Parameters ---------- @@ -110,8 +106,6 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): The timezeone to use for the data query cutoff. dataset: DataSet The DataSet object for which this loader loads data. - concrete_loader: EventsLoader - The reference loader to use for this dataset. Notes ----- @@ -145,6 +139,8 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): SHARE_COUNT_FIELD_NAME, }) + concrete_loader = ShareBuybackAuthorizationsLoader + def __init__(self, expr, resources=None, @@ -162,7 +158,3 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): data_query_tz=data_query_tz, dataset=dataset, **kwargs) - - @property - def concrete_loader(self): - return ShareBuybackAuthorizationsLoader diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 80b963d0..39724e76 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -57,6 +57,8 @@ class BlazeEarningsCalendarLoader(BlazeEventsLoader): ANNOUNCEMENT_FIELD_NAME, }) + concrete_loader = EarningsCalendarLoader + def __init__(self, expr, resources=None, @@ -71,7 +73,3 @@ class BlazeEarningsCalendarLoader(BlazeEventsLoader): resources=resources, odo_kwargs=odo_kwargs, data_query_time=data_query_time, data_query_tz=data_query_tz, **kwargs) - - @property - def concrete_loader(self): - return EarningsCalendarLoader diff --git a/zipline/pipeline/loaders/blaze/events.py b/zipline/pipeline/loaders/blaze/events.py index 3d334ce6..93ed6b72 100644 --- a/zipline/pipeline/loaders/blaze/events.py +++ b/zipline/pipeline/loaders/blaze/events.py @@ -124,4 +124,5 @@ class BlazeEventsLoader(PipelineLoader): ).load_adjusted_array(columns, dates, assets, mask) def prepare_data(self, raw, gb): - return {sid: raw.loc[group] for sid, group in gb.groups.items()} + return {sid: raw.loc[group].drop(SID_FIELD_NAME, axis=1) for sid, group + in gb.groups.items()} diff --git a/zipline/pipeline/loaders/buyback_auth.py b/zipline/pipeline/loaders/buyback_auth.py index 677732ec..d46a207a 100644 --- a/zipline/pipeline/loaders/buyback_auth.py +++ b/zipline/pipeline/loaders/buyback_auth.py @@ -24,6 +24,8 @@ class CashBuybackAuthorizationsLoader(EventsLoader): event date, cash value)] """ + expected_cols = frozenset([BUYBACK_ANNOUNCEMENT_FIELD_NAME, + CASH_FIELD_NAME]) def __init__(self, all_dates, @@ -37,10 +39,6 @@ class CashBuybackAuthorizationsLoader(EventsLoader): dataset=dataset, ) - @property - def expected_cols(self): - return frozenset([BUYBACK_ANNOUNCEMENT_FIELD_NAME, CASH_FIELD_NAME]) - @lazyval def previous_value_loader(self): return self._previous_event_value_loader( @@ -69,6 +67,8 @@ class ShareBuybackAuthorizationsLoader(EventsLoader): event date, share value)] """ + expected_cols = frozenset([BUYBACK_ANNOUNCEMENT_FIELD_NAME, + SHARE_COUNT_FIELD_NAME]) def __init__(self, all_dates, @@ -82,11 +82,6 @@ class ShareBuybackAuthorizationsLoader(EventsLoader): dataset=dataset, ) - @property - def expected_cols(self): - return frozenset([BUYBACK_ANNOUNCEMENT_FIELD_NAME, - SHARE_COUNT_FIELD_NAME]) - @lazyval def previous_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 f3225790..89d376c9 100644 --- a/zipline/pipeline/loaders/earnings.py +++ b/zipline/pipeline/loaders/earnings.py @@ -10,6 +10,8 @@ from zipline.utils.memoize import lazyval class EarningsCalendarLoader(EventsLoader): + expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME]) + def __init__(self, all_dates, events_by_sid, infer_timestamps=False, dataset=EarningsCalendar): @@ -17,10 +19,6 @@ class EarningsCalendarLoader(EventsLoader): all_dates, events_by_sid, infer_timestamps, dataset=dataset, ) - @property - def expected_cols(self): - return frozenset([ANNOUNCEMENT_FIELD_NAME]) - @lazyval def next_announcement_loader(self): return self._next_event_date_loader(self.dataset.next_announcement, diff --git a/zipline/pipeline/loaders/events.py b/zipline/pipeline/loaders/events.py index dd20b7c5..84d0735c 100644 --- a/zipline/pipeline/loaders/events.py +++ b/zipline/pipeline/loaders/events.py @@ -6,13 +6,18 @@ from toolz import merge from .base import PipelineLoader from .frame import DataFrameLoader from .utils import next_date_frame, previous_date_frame, previous_value -from zipline.pipeline.common import SID_FIELD_NAME, TS_FIELD_NAME +from zipline.pipeline.common import TS_FIELD_NAME WRONG_COLS_ERROR = "Expected columns {expected_columns} for sid {sid} but " \ "got columns {resulting_columns}." -BAD_DATA_FORMAT_ERROR = ("Data for sid {sid} must be in DataFrame, " - "Series, or DatetimeIndex.") +WRONG_SINGLE_COL_DATA_FORMAT_ERROR = ("Data for sid {sid} is expected to have " + "1 column and to be in a DataFrame, " + "Series, or DatetimeIndex.") + +WRONG_MANY_COL_DATA_FORMAT_ERROR = ("Data for sid {sid} is expected to have " + "more than 1 column and to be in a " + "DataFrame.") SERIES_NO_DTINDEX_ERROR = ("Got Series for sid {sid}, but index was not " "DatetimeIndex.") @@ -85,23 +90,8 @@ class EventsLoader(PipelineLoader): dates = self.all_dates.values for k, v in iteritems(events_by_sid): - # First, must convert to DataFrame. - if isinstance(v, pd.Series): - if not isinstance(v.index, pd.DatetimeIndex): - raise ValueError( - SERIES_NO_DTINDEX_ERROR.format(sid=k) - ) - self.events_by_sid[k] = v = pd.DataFrame(v) - elif isinstance(v, pd.DatetimeIndex): - if not infer_timestamps: - raise ValueError( - DTINDEX_NOT_INFER_TS_ERROR.format(sid=k) - ) - self.events_by_sid[k] = v = pd.DataFrame( - v, index=[dates[0]] * len(v) - ) # Already a DataFrame - elif isinstance(v, pd.DataFrame): + if isinstance(v, pd.DataFrame): if TS_FIELD_NAME not in v.columns: if not infer_timestamps: raise ValueError( @@ -114,21 +104,50 @@ class EventsLoader(PipelineLoader): v.index = [dates[0]] * len(v) else: self.events_by_sid[k] = v.set_index(TS_FIELD_NAME) - else: - raise ValueError(BAD_DATA_FORMAT_ERROR.format(sid=k)) - # Once data is in a DF, make sure columns are correct. - cols_except_ts = (set(v.columns) - - {TS_FIELD_NAME} - - {SID_FIELD_NAME}) - # Check that all columns other than timestamp are as expected. - if cols_except_ts != self.expected_cols: - raise ValueError( - WRONG_COLS_ERROR .format( - expected_columns=self.expected_cols, - sid=k, - resulting_columns=v.columns.values + # Once data is in a DF, make sure columns are correct. + cols_except_ts = (set(v.columns) - + {TS_FIELD_NAME}) + + # Check that all columns other than timestamp are as expected. + if cols_except_ts != self.expected_cols: + raise ValueError( + WRONG_COLS_ERROR.format( + expected_columns=list(self.expected_cols), + sid=k, + resulting_columns=v.columns.values + ) ) + # Not a DataFrame and we only expect 1 column + elif len(self.expected_cols) == 1: + # First, must convert to DataFrame. + if isinstance(v, pd.Series): + if not isinstance(v.index, pd.DatetimeIndex): + raise ValueError( + SERIES_NO_DTINDEX_ERROR.format(sid=k) + ) + self.events_by_sid[k] = pd.DataFrame({ + list(self.expected_cols)[0]: v}) + elif isinstance(v, pd.DatetimeIndex): + if not infer_timestamps: + raise ValueError( + DTINDEX_NOT_INFER_TS_ERROR.format(sid=k) + ) + self.events_by_sid[k] = pd.DataFrame({ + list(self.expected_cols)[0]: v + }, index=[dates[0]] * len(v)) + else: + # We expect 1 column, but we got something other than a + # Series, DatetimeIndex, or DataFrame. + raise ValueError( + WRONG_SINGLE_COL_DATA_FORMAT_ERROR.format(sid=k) + ) + else: + # We expected multiple columns, but we got something other + # than a DataFrame. + raise ValueError( + WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=k) ) + self.dataset = dataset def get_loader(self, column):