mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-18 12:20:12 +08:00
237 lines
8.9 KiB
Python
237 lines
8.9 KiB
Python
"""
|
|
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
|
|
|
|
from zipline.pipeline.common import (
|
|
ANNOUNCEMENT_FIELD_NAME,
|
|
DAYS_SINCE_PREV,
|
|
DAYS_TO_NEXT,
|
|
NEXT_ANNOUNCEMENT,
|
|
PREVIOUS_ANNOUNCEMENT,
|
|
SID_FIELD_NAME,
|
|
TS_FIELD_NAME
|
|
)
|
|
from zipline.pipeline.data import EarningsCalendar
|
|
from zipline.pipeline.factors.events import (
|
|
BusinessDaysSincePreviousEarnings,
|
|
BusinessDaysUntilNextEarnings,
|
|
)
|
|
from zipline.pipeline.loaders.earnings import EarningsCalendarLoader
|
|
from zipline.pipeline.loaders.blaze import (
|
|
BlazeEarningsCalendarLoader,
|
|
)
|
|
|
|
from zipline.testing import tmp_asset_finder
|
|
|
|
earnings_cases = [
|
|
# K1--K2--A1--A2.
|
|
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--A2--A1.
|
|
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--A1--K2--A2.
|
|
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(
|
|
columns=[ANNOUNCEMENT_FIELD_NAME,
|
|
TS_FIELD_NAME],
|
|
dtype='datetime64[ns]'
|
|
),
|
|
]
|
|
|
|
|
|
class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin):
|
|
"""
|
|
Tests for loading the earnings announcement data.
|
|
"""
|
|
pipeline_columns = {
|
|
NEXT_ANNOUNCEMENT: EarningsCalendar.next_announcement.latest,
|
|
PREVIOUS_ANNOUNCEMENT: EarningsCalendar.previous_announcement.latest,
|
|
DAYS_SINCE_PREV: BusinessDaysSincePreviousEarnings(),
|
|
DAYS_TO_NEXT: BusinessDaysUntilNextEarnings(),
|
|
}
|
|
|
|
@classmethod
|
|
def get_sids(cls):
|
|
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()),
|
|
)
|
|
|
|
cls.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)),
|
|
}, 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)),
|
|
}, index=dates)
|
|
|
|
@classmethod
|
|
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_next_busday_offsets = self._compute_busday_offsets(
|
|
_expected_next_announce
|
|
)
|
|
_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
|
|
|
|
|
|
class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super(BlazeEarningsCalendarLoaderTestCase, cls).setUpClass()
|
|
cls.loader_type = BlazeEarningsCalendarLoader
|
|
|
|
def loader_args(self, dates):
|
|
_, mapping = super(
|
|
BlazeEarningsCalendarLoaderTestCase,
|
|
self,
|
|
).loader_args(dates)
|
|
return (bz.data(pd.concat(
|
|
pd.DataFrame({
|
|
ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME],
|
|
TS_FIELD_NAME: df[TS_FIELD_NAME],
|
|
SID_FIELD_NAME: sid,
|
|
})
|
|
for sid, df in iteritems(mapping)
|
|
).reset_index(drop=True)),)
|
|
|
|
|
|
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(
|
|
BlazeEarningsCalendarLoaderNotInteractiveTestCase,
|
|
self,
|
|
).loader_args(dates)
|
|
return swap_resources_into_scope(bound_expr, {})
|