wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-31 12:10:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

MAINT: modify coersion logic.

Browse Source 
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.
This commit is contained in:
Maya Tydykov
2016-02-26 09:51:58 -05:00
parent 5b37af6e04
commit 18d838be87
9 changed files with 212 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/test_buyback_auth.py
+26 -37
View File
@@ -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.
tests/pipeline/test_earnings.py
+6 -39
View File
@@ -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
tests/pipeline/test_events.py
+111 -45
View File
@@ -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)
zipline/pipeline/loaders/blaze/buyback_auth.py
+8 -16
View File
@@ -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
zipline/pipeline/loaders/blaze/earnings.py
+2 -4
View File
@@ -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
zipline/pipeline/loaders/blaze/events.py
+2 -1
View File
@@ -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()}
zipline/pipeline/loaders/buyback_auth.py
+4 -9
View File
@@ -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(
zipline/pipeline/loaders/earnings.py
+2 -4
View File
@@ -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,
zipline/pipeline/loaders/events.py
+51 -32
View File
@@ -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):