mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-12 10:39:22 +08:00
TST: enhance test; add more common data.
STY: fixes for PEP8.
This commit is contained in:
@@ -8,13 +8,10 @@ import blaze as bz
|
||||
from blaze.compute.core import swap_resources_into_scope
|
||||
from contextlib2 import ExitStack
|
||||
from nose_parameterized import parameterized
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from pandas.util.testing import assert_series_equal
|
||||
from six import iteritems
|
||||
from tests.pipeline.test_events import param_dates, EventLoaderCommonTest
|
||||
|
||||
from zipline.pipeline import Pipeline
|
||||
from zipline.pipeline.common import(
|
||||
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
|
||||
CASH_FIELD_NAME,
|
||||
@@ -27,7 +24,6 @@ from zipline.pipeline.common import(
|
||||
TS_FIELD_NAME)
|
||||
from zipline.pipeline.data import (CashBuybackAuthorizations,
|
||||
ShareBuybackAuthorizations)
|
||||
from zipline.pipeline.engine import SimplePipelineEngine
|
||||
from zipline.pipeline.factors.events import (
|
||||
BusinessDaysSincePreviousCashBuybackAuth,
|
||||
BusinessDaysSincePreviousShareBuybackAuth
|
||||
@@ -38,11 +34,7 @@ from zipline.pipeline.loaders.blaze import (
|
||||
BlazeCashBuybackAuthorizationsLoader,
|
||||
BlazeShareBuybackAuthorizationsLoader,
|
||||
)
|
||||
from zipline.utils.numpy_utils import make_datetime64D, NaTD
|
||||
from zipline.utils.test_utils import (
|
||||
gen_calendars,
|
||||
make_simple_equity_info,
|
||||
num_days_in_range,
|
||||
tmp_asset_finder,
|
||||
)
|
||||
|
||||
@@ -97,7 +89,7 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest):
|
||||
Test for cash buyback authorizations dataset.
|
||||
"""
|
||||
pipeline_columns = {
|
||||
('%s' % PREVIOUS_BUYBACK_CASH):
|
||||
(PREVIOUS_BUYBACK_CASH):
|
||||
CashBuybackAuthorizations.previous_value.latest,
|
||||
PREVIOUS_BUYBACK_ANNOUNCEMENT:
|
||||
CashBuybackAuthorizations.previous_announcement_date.latest,
|
||||
@@ -124,8 +116,6 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest):
|
||||
zip_with_floats_dates = partial(self.zip_with_floats, dates)
|
||||
num_days_between_dates = partial(self.num_days_between, dates)
|
||||
_expected_previous_cash = pd.DataFrame({
|
||||
# TODO if the next knowledge date is 10, why is the range
|
||||
# until 15?
|
||||
0: zip_with_floats_dates(
|
||||
['NaN'] * num_days_between_dates(None, '2014-01-14') +
|
||||
[10] * num_days_between_dates('2014-01-15', '2014-01-19') +
|
||||
@@ -148,14 +138,16 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest):
|
||||
),
|
||||
4: zip_with_floats_dates(['NaN'] * len(dates)),
|
||||
}, index=dates)
|
||||
self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] = self.get_expected_previous(
|
||||
dates)
|
||||
self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] = \
|
||||
self.get_expected_previous_event_dates(dates)
|
||||
self.cols[PREVIOUS_BUYBACK_CASH] = _expected_previous_cash
|
||||
|
||||
self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
|
||||
self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT]
|
||||
)
|
||||
|
||||
@parameterized.expand(param_dates)
|
||||
def test_compute_cash_buyback_auth(self, dates):
|
||||
self._test_compute_buyback_auth(dates)
|
||||
self._test_compute(dates)
|
||||
|
||||
|
||||
class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase):
|
||||
@@ -163,9 +155,9 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase):
|
||||
Test for share buyback authorizations dataset.
|
||||
"""
|
||||
pipeline_columns = {
|
||||
('%s' % PREVIOUS_BUYBACK_SHARE_COUNT):
|
||||
PREVIOUS_BUYBACK_SHARE_COUNT:
|
||||
ShareBuybackAuthorizations.previous_share_count.latest,
|
||||
('%s' % PREVIOUS_BUYBACK_ANNOUNCEMENT):
|
||||
PREVIOUS_BUYBACK_ANNOUNCEMENT:
|
||||
ShareBuybackAuthorizations.previous_announcement_date.latest,
|
||||
DAYS_SINCE_PREV:
|
||||
BusinessDaysSincePreviousShareBuybackAuth(),
|
||||
@@ -179,8 +171,8 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase):
|
||||
)
|
||||
cls.cols = {}
|
||||
cls.dataset = {sid: df.drop(CASH_FIELD_NAME, 1)
|
||||
for sid, df in
|
||||
enumerate(buyback_authorizations)}
|
||||
for sid, df in
|
||||
enumerate(buyback_authorizations)}
|
||||
cls.loader_type = ShareBuybackAuthorizationsLoader
|
||||
|
||||
@classmethod
|
||||
@@ -217,11 +209,14 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase):
|
||||
PREVIOUS_BUYBACK_SHARE_COUNT
|
||||
] = _expected_previous_buyback_share_count
|
||||
self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] = \
|
||||
self.get_expected_previous(dates)
|
||||
self.get_expected_previous_event_dates(dates)
|
||||
self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
|
||||
self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT]
|
||||
)
|
||||
|
||||
@parameterized.expand(param_dates)
|
||||
def test_compute_share_buyback_auth(self, dates):
|
||||
self._test_compute_buyback_auth(dates)
|
||||
self._test_compute(dates)
|
||||
|
||||
|
||||
class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase):
|
||||
|
||||
@@ -35,35 +35,35 @@ from zipline.utils.test_utils import (
|
||||
)
|
||||
|
||||
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([])
|
||||
})
|
||||
]
|
||||
# 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):
|
||||
@@ -97,11 +97,12 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonTest):
|
||||
def tearDownClass(cls):
|
||||
cls._cleanup_stack.close()
|
||||
|
||||
|
||||
def setup(self, dates):
|
||||
_expected_next_announce = self.get_expected_next_event_dates(dates)
|
||||
|
||||
_expected_previous_announce = self.get_expected_previous_event_dates(dates)
|
||||
_expected_previous_announce = self.get_expected_previous_event_dates(
|
||||
dates
|
||||
)
|
||||
|
||||
_expected_next_busday_offsets = self._compute_busday_offsets(
|
||||
_expected_next_announce
|
||||
@@ -146,7 +147,8 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase(
|
||||
"""
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super(BlazeEarningsCalendarLoaderNotInteractiveTestCase, cls).setUpClass()
|
||||
super(BlazeEarningsCalendarLoaderNotInteractiveTestCase,
|
||||
cls).setUpClass()
|
||||
cls.loader_type = BlazeEarningsCalendarLoader
|
||||
|
||||
def loader_args(self, dates):
|
||||
|
||||
@@ -26,7 +26,9 @@ from zipline.pipeline.loaders.events import (
|
||||
WRONG_COLS_ERROR,
|
||||
)
|
||||
from zipline.utils.memoize import lazyval
|
||||
from zipline.utils.numpy_utils import datetime64ns_dtype, NaTD, make_datetime64D
|
||||
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
|
||||
|
||||
@@ -153,25 +155,42 @@ class EventLoaderTestCase(TestCase):
|
||||
expected,
|
||||
)
|
||||
|
||||
@parameterized.expand([
|
||||
# DataFrame without timestamp column and infer_timestamps = True
|
||||
[pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}),
|
||||
False,
|
||||
DF_NO_TS_NOT_INFER_TS_ERROR.format(
|
||||
timestamp_column_name=TS_FIELD_NAME,
|
||||
sid=0
|
||||
)
|
||||
],
|
||||
# DatetimeIndex with infer_timestamps = False
|
||||
[pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME), False,
|
||||
DTINDEX_NOT_INFER_TS_ERROR.format(sid=0)],
|
||||
# Series with DatetimeIndex as index and infer_timestamps = False
|
||||
[pd.Series(dtx, name=ANNOUNCEMENT_FIELD_NAME), False,
|
||||
SERIES_NO_DTINDEX_ERROR.format(sid=0)],
|
||||
# Some other data structure that is not expected
|
||||
[dtx, False, BAD_DATA_FORMAT_ERROR.format(sid=0)],
|
||||
[dtx, True, BAD_DATA_FORMAT_ERROR.format(sid=0)]
|
||||
])
|
||||
@parameterized.expand(
|
||||
[
|
||||
# DataFrame without timestamp column and infer_timestamps = True
|
||||
[
|
||||
pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}),
|
||||
False,
|
||||
DF_NO_TS_NOT_INFER_TS_ERROR.format(
|
||||
timestamp_column_name=TS_FIELD_NAME,
|
||||
sid=0
|
||||
)
|
||||
],
|
||||
# DatetimeIndex with infer_timestamps = False
|
||||
[
|
||||
pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME),
|
||||
False,
|
||||
DTINDEX_NOT_INFER_TS_ERROR.format(sid=0)
|
||||
],
|
||||
# Series with DatetimeIndex as index and infer_timestamps = False
|
||||
[
|
||||
pd.Series(dtx, name=ANNOUNCEMENT_FIELD_NAME),
|
||||
False,
|
||||
SERIES_NO_DTINDEX_ERROR.format(sid=0)
|
||||
],
|
||||
# Some other data structure that is not expected
|
||||
[
|
||||
dtx,
|
||||
False,
|
||||
BAD_DATA_FORMAT_ERROR.format(sid=0)
|
||||
],
|
||||
[
|
||||
dtx,
|
||||
True,
|
||||
BAD_DATA_FORMAT_ERROR.format(sid=0)
|
||||
]
|
||||
]
|
||||
)
|
||||
def test_bad_conversion_to_df(self, df, infer_timestamps, msg):
|
||||
events_by_sid = {0: df}
|
||||
assert_loader_error(events_by_sid, ValueError, msg,
|
||||
@@ -204,14 +223,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
|
||||
# once.
|
||||
param_dates = list(gen_calendars(
|
||||
@@ -237,15 +248,12 @@ class EventLoaderCommonTest(object):
|
||||
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)
|
||||
|
||||
|
||||
def loader_args(self, dates):
|
||||
"""Construct the base object to pass to the loader.
|
||||
|
||||
@@ -268,7 +276,59 @@ class EventLoaderCommonTest(object):
|
||||
loader = self.loader_type(*self.loader_args(dates))
|
||||
return SimplePipelineEngine(lambda _: loader, dates, self.finder)
|
||||
|
||||
def get_expected_previous(self, dates):
|
||||
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)),
|
||||
}, 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({
|
||||
@@ -339,7 +399,7 @@ class EventLoaderCommonTest(object):
|
||||
index=announcement_dates.index,
|
||||
)
|
||||
|
||||
def _test_compute_buyback_auth(self, dates):
|
||||
def _test_compute(self, dates):
|
||||
engine = self.setup_engine(dates)
|
||||
self.setup(dates)
|
||||
|
||||
@@ -358,7 +418,3 @@ class EventLoaderCommonTest(object):
|
||||
assert_series_equal(result[col_name].xs(sid, level=1),
|
||||
self.cols[col_name][sid],
|
||||
check_names=False)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user