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

TST: recfator tests to use fixtures

Browse Source 
MAINT: use np.array

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