mirror of
https://github.com/wassname/catalyst.git
synced 2026-09-12 12:12:04 +08:00
ENH: add custom dataset for estimize
MAINT: alphabetize constants MAINT: remove obsolete column TST: refactor tests to use common code MAINT: remove unneeded fields from dataset MAINT: remove obsolete earnings estimates columns and refactor
This commit is contained in:
Maya Tydykov
@@ -34,7 +34,6 @@ from zipline.pipeline.loaders.blaze import (
|
||||
BlazeShareBuybackAuthorizationsLoader,
|
||||
)
|
||||
from zipline.pipeline.loaders.utils import (
|
||||
get_values_for_date_ranges,
|
||||
zip_with_floats,
|
||||
zip_with_dates
|
||||
)
|
||||
@@ -42,8 +41,9 @@ from zipline.testing.fixtures import (
|
||||
WithPipelineEventDataLoader, ZiplineTestCase
|
||||
)
|
||||
|
||||
date_intervals = [[None, '2014-01-04'], ['2014-01-05', '2014-01-09'],
|
||||
['2014-01-10', None]]
|
||||
date_intervals = [
|
||||
[[None, '2014-01-04'], ['2014-01-05', '2014-01-09'], ['2014-01-10', None]]
|
||||
]
|
||||
|
||||
buyback_authorizations_cases = [
|
||||
pd.DataFrame({
|
||||
@@ -63,18 +63,6 @@ buyback_authorizations_cases = [
|
||||
]
|
||||
|
||||
|
||||
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_date_index_with_vals,
|
||||
vals_for_date_intervals,
|
||||
date_intervals,
|
||||
dates),
|
||||
1: zip_date_index_with_vals(dates, ['NaN'] * len(dates)),
|
||||
}, index=dates)
|
||||
|
||||
|
||||
class CashBuybackAuthLoaderTestCase(WithPipelineEventDataLoader,
|
||||
ZiplineTestCase):
|
||||
"""
|
||||
@@ -103,16 +91,17 @@ class CashBuybackAuthLoaderTestCase(WithPipelineEventDataLoader,
|
||||
loader_type = CashBuybackAuthorizationsLoader
|
||||
|
||||
def setup(self, dates):
|
||||
cols = {}
|
||||
_expected_previous_cash = get_expected_previous_values(
|
||||
zip_with_floats, dates,
|
||||
['NaN', 10, 20]
|
||||
)
|
||||
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 = {
|
||||
PREVIOUS_BUYBACK_CASH: self.get_sids_to_frames(zip_with_floats,
|
||||
[['NaN', 10, 20]],
|
||||
date_intervals,
|
||||
dates),
|
||||
PREVIOUS_BUYBACK_ANNOUNCEMENT: self.get_sids_to_frames(
|
||||
zip_with_dates,
|
||||
[['NaT', '2014-01-04', '2014-01-09']],
|
||||
date_intervals,
|
||||
dates),
|
||||
}
|
||||
cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
|
||||
cols[PREVIOUS_BUYBACK_ANNOUNCEMENT]
|
||||
)
|
||||
@@ -147,16 +136,18 @@ class ShareBuybackAuthLoaderTestCase(WithPipelineEventDataLoader,
|
||||
loader_type = ShareBuybackAuthorizationsLoader
|
||||
|
||||
def setup(self, dates):
|
||||
cols = {}
|
||||
cols[
|
||||
PREVIOUS_BUYBACK_SHARE_COUNT
|
||||
] = get_expected_previous_values(zip_with_floats,
|
||||
dates,
|
||||
['NaN', 1, 15])
|
||||
cols[
|
||||
PREVIOUS_BUYBACK_ANNOUNCEMENT
|
||||
] = get_expected_previous_values(zip_with_dates, dates,
|
||||
['NaT', '2014-01-04', '2014-01-09'])
|
||||
cols = {PREVIOUS_BUYBACK_SHARE_COUNT:
|
||||
self.get_sids_to_frames(zip_with_floats,
|
||||
[['NaN', 1, 15]],
|
||||
date_intervals,
|
||||
dates,),
|
||||
PREVIOUS_BUYBACK_ANNOUNCEMENT:
|
||||
self.get_sids_to_frames(zip_with_dates,
|
||||
[['NaT', '2014-01-04', '2014-01-09']],
|
||||
date_intervals,
|
||||
dates,),
|
||||
}
|
||||
|
||||
cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
|
||||
cols[PREVIOUS_BUYBACK_ANNOUNCEMENT]
|
||||
)
|
||||
|
||||
@@ -0,0 +1,327 @@
|
||||
"""
|
||||
Tests for the reference loader for ConsensusEstimates.
|
||||
"""
|
||||
import blaze as bz
|
||||
from blaze.compute.core import swap_resources_into_scope
|
||||
import pandas as pd
|
||||
from six import iteritems
|
||||
|
||||
from zipline.pipeline.common import (
|
||||
COUNT_FIELD_NAME,
|
||||
FISCAL_QUARTER_FIELD_NAME,
|
||||
FISCAL_YEAR_FIELD_NAME,
|
||||
HIGH_FIELD_NAME,
|
||||
LOW_FIELD_NAME,
|
||||
MEAN_FIELD_NAME,
|
||||
NEXT_COUNT,
|
||||
NEXT_FISCAL_QUARTER,
|
||||
NEXT_FISCAL_YEAR,
|
||||
NEXT_HIGH,
|
||||
NEXT_LOW,
|
||||
NEXT_RELEASE_DATE,
|
||||
NEXT_STANDARD_DEVIATION,
|
||||
PREVIOUS_COUNT,
|
||||
PREVIOUS_FISCAL_QUARTER,
|
||||
PREVIOUS_FISCAL_YEAR,
|
||||
PREVIOUS_HIGH,
|
||||
PREVIOUS_LOW,
|
||||
PREVIOUS_MEAN, NEXT_MEAN,
|
||||
PREVIOUS_RELEASE_DATE,
|
||||
PREVIOUS_STANDARD_DEVIATION,
|
||||
RELEASE_DATE_FIELD_NAME,
|
||||
STANDARD_DEVIATION_FIELD_NAME,
|
||||
SID_FIELD_NAME)
|
||||
from zipline.pipeline.data import ConsensusEstimates
|
||||
from zipline.pipeline.loaders.consensus_estimates import (
|
||||
ConsensusEstimatesLoader
|
||||
)
|
||||
from zipline.pipeline.loaders.blaze import BlazeConsensusEstimatesLoader
|
||||
from zipline.pipeline.loaders.utils import (
|
||||
zip_with_floats
|
||||
)
|
||||
from zipline.testing.fixtures import (
|
||||
ZiplineTestCase,
|
||||
WithNextAndPreviousEventDataLoader
|
||||
)
|
||||
|
||||
consensus_estimates_cases = [
|
||||
# K1--K2--A1--A2.
|
||||
pd.DataFrame({
|
||||
STANDARD_DEVIATION_FIELD_NAME: (.5, .6),
|
||||
COUNT_FIELD_NAME: (1, 2),
|
||||
FISCAL_QUARTER_FIELD_NAME: (1, 1),
|
||||
HIGH_FIELD_NAME: (.6, .7),
|
||||
MEAN_FIELD_NAME: (.1, .2),
|
||||
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
|
||||
LOW_FIELD_NAME: (.05, .06),
|
||||
}),
|
||||
# K1--K2--A2--A1.
|
||||
pd.DataFrame({
|
||||
STANDARD_DEVIATION_FIELD_NAME: (.6, .7),
|
||||
COUNT_FIELD_NAME: (2, 3),
|
||||
FISCAL_QUARTER_FIELD_NAME: (1, 1),
|
||||
HIGH_FIELD_NAME: (.7, .8),
|
||||
MEAN_FIELD_NAME: (.2, .3),
|
||||
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
|
||||
LOW_FIELD_NAME: (.06, .07),
|
||||
}),
|
||||
# K1--A1--K2--A2.
|
||||
pd.DataFrame({
|
||||
STANDARD_DEVIATION_FIELD_NAME: (.7, .8),
|
||||
COUNT_FIELD_NAME: (3, 4),
|
||||
FISCAL_QUARTER_FIELD_NAME: (1, 1),
|
||||
HIGH_FIELD_NAME: (.8, .9),
|
||||
MEAN_FIELD_NAME: (.3, .4),
|
||||
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
|
||||
LOW_FIELD_NAME: (.07, .08),
|
||||
}),
|
||||
# K1 == K2.
|
||||
pd.DataFrame({
|
||||
STANDARD_DEVIATION_FIELD_NAME: (.8, .9),
|
||||
COUNT_FIELD_NAME: (4, 5),
|
||||
FISCAL_QUARTER_FIELD_NAME: (1, 1),
|
||||
HIGH_FIELD_NAME: (.9, 1.0),
|
||||
MEAN_FIELD_NAME: (.4, .5),
|
||||
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
|
||||
LOW_FIELD_NAME: (.08, .09),
|
||||
}),
|
||||
pd.DataFrame(
|
||||
columns=[STANDARD_DEVIATION_FIELD_NAME,
|
||||
COUNT_FIELD_NAME,
|
||||
FISCAL_QUARTER_FIELD_NAME,
|
||||
HIGH_FIELD_NAME,
|
||||
MEAN_FIELD_NAME,
|
||||
FISCAL_YEAR_FIELD_NAME,
|
||||
LOW_FIELD_NAME],
|
||||
dtype='datetime64[ns]'
|
||||
),
|
||||
]
|
||||
|
||||
next_standard_deviation = [
|
||||
['NaN', .5, .6, 'NaN'],
|
||||
['NaN', .6, .7, .6, 'NaN'],
|
||||
['NaN', .7, 'NaN', .8, 'NaN'],
|
||||
['NaN', .8, .9, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_standard_deviation = [
|
||||
['NaN', .5, .6],
|
||||
['NaN', .7, .6],
|
||||
['NaN', .7, .8],
|
||||
['NaN', .8, .9],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
next_count = [
|
||||
['NaN', 1, 2, 'NaN'],
|
||||
['NaN', 2, 3, 2, 'NaN'],
|
||||
['NaN', 3, 'NaN', 4, 'NaN'],
|
||||
['NaN', 4, 5, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_count = [
|
||||
['NaN', 1, 2],
|
||||
['NaN', 3, 2],
|
||||
['NaN', 3, 4],
|
||||
['NaN', 4, 5],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
next_fiscal_quarter = [
|
||||
['NaN', 1, 1, 'NaN'],
|
||||
['NaN', 1, 1, 1, 'NaN'],
|
||||
['NaN', 1, 'NaN', 1, 'NaN'],
|
||||
['NaN', 1, 1, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_fiscal_quarter = [
|
||||
['NaN', 1, 1],
|
||||
['NaN', 1, 1],
|
||||
['NaN', 1, 1],
|
||||
['NaN', 1, 1],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
next_high = [
|
||||
['NaN', .6, .7, 'NaN'],
|
||||
['NaN', .7, .8, .7, 'NaN'],
|
||||
['NaN', .8, 'NaN', .9, 'NaN'],
|
||||
['NaN', .9, 1.0, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_high = [
|
||||
['NaN', .6, .7],
|
||||
['NaN', .8, .7],
|
||||
['NaN', .8, .9],
|
||||
['NaN', .9, 1.0],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
next_mean = [
|
||||
['NaN', .1, .2, 'NaN'],
|
||||
['NaN', .2, .3, .2, 'NaN'],
|
||||
['NaN', .3, 'NaN', .4, 'NaN'],
|
||||
['NaN', .4, .5, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_mean = [
|
||||
['NaN', .1, .2],
|
||||
['NaN', .3, .2],
|
||||
['NaN', .3, .4],
|
||||
['NaN', .4, .5],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
next_fiscal_year = [
|
||||
['NaN', 2014, 2014, 'NaN'],
|
||||
['NaN', 2014, 2014, 2014, 'NaN'],
|
||||
['NaN', 2014, 'NaN', 2014, 'NaN'],
|
||||
['NaN', 2014, 2014, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_fiscal_year = [
|
||||
['NaN', 2014, 2014],
|
||||
['NaN', 2014, 2014],
|
||||
['NaN', 2014, 2014],
|
||||
['NaN', 2014, 2014],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
next_low = [
|
||||
['NaN', .05, .06, 'NaN'],
|
||||
['NaN', .06, .07, .06, 'NaN'],
|
||||
['NaN', .07, 'NaN', .08, 'NaN'],
|
||||
['NaN', .08, .09, 'NaN'],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
prev_low = [
|
||||
['NaN', .05, .06],
|
||||
['NaN', .07, .06],
|
||||
['NaN', .07, .08],
|
||||
['NaN', .08, .09],
|
||||
['NaN']
|
||||
]
|
||||
|
||||
field_name_to_expected_col = {
|
||||
PREVIOUS_STANDARD_DEVIATION: prev_standard_deviation,
|
||||
NEXT_STANDARD_DEVIATION: next_standard_deviation,
|
||||
PREVIOUS_COUNT: prev_count,
|
||||
NEXT_COUNT: next_count,
|
||||
PREVIOUS_FISCAL_QUARTER: prev_fiscal_quarter,
|
||||
NEXT_FISCAL_QUARTER: next_fiscal_quarter,
|
||||
PREVIOUS_HIGH: prev_high,
|
||||
NEXT_HIGH: next_high,
|
||||
PREVIOUS_MEAN: prev_mean,
|
||||
NEXT_MEAN: next_mean,
|
||||
PREVIOUS_FISCAL_YEAR: prev_fiscal_year,
|
||||
NEXT_FISCAL_YEAR: next_fiscal_year,
|
||||
PREVIOUS_LOW: prev_low,
|
||||
NEXT_LOW: next_low
|
||||
}
|
||||
|
||||
|
||||
class ConsensusEstimatesLoaderTestCase(WithNextAndPreviousEventDataLoader,
|
||||
ZiplineTestCase):
|
||||
"""
|
||||
Tests for loading the consensus estimates data.
|
||||
"""
|
||||
pipeline_columns = {
|
||||
NEXT_RELEASE_DATE:
|
||||
ConsensusEstimates.next_release_date.latest,
|
||||
PREVIOUS_RELEASE_DATE:
|
||||
ConsensusEstimates.previous_release_date.latest,
|
||||
PREVIOUS_STANDARD_DEVIATION:
|
||||
ConsensusEstimates.previous_standard_deviation.latest,
|
||||
NEXT_STANDARD_DEVIATION:
|
||||
ConsensusEstimates.next_standard_deviation.latest,
|
||||
PREVIOUS_COUNT:
|
||||
ConsensusEstimates.previous_count.latest,
|
||||
NEXT_COUNT:
|
||||
ConsensusEstimates.next_count.latest,
|
||||
PREVIOUS_FISCAL_QUARTER:
|
||||
ConsensusEstimates.previous_fiscal_quarter.latest,
|
||||
NEXT_FISCAL_QUARTER:
|
||||
ConsensusEstimates.next_fiscal_quarter.latest,
|
||||
PREVIOUS_HIGH:
|
||||
ConsensusEstimates.previous_high.latest,
|
||||
NEXT_HIGH:
|
||||
ConsensusEstimates.next_high.latest,
|
||||
PREVIOUS_MEAN:
|
||||
ConsensusEstimates.previous_mean.latest,
|
||||
NEXT_MEAN:
|
||||
ConsensusEstimates.next_mean.latest,
|
||||
PREVIOUS_FISCAL_YEAR:
|
||||
ConsensusEstimates.previous_fiscal_year.latest,
|
||||
NEXT_FISCAL_YEAR:
|
||||
ConsensusEstimates.next_fiscal_year.latest,
|
||||
PREVIOUS_LOW:
|
||||
ConsensusEstimates.previous_low.latest,
|
||||
NEXT_LOW:
|
||||
ConsensusEstimates.next_low.latest
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_dataset(cls):
|
||||
return {sid:
|
||||
pd.concat([
|
||||
cls.base_cases[sid].rename(columns={
|
||||
'other_date': RELEASE_DATE_FIELD_NAME
|
||||
}),
|
||||
df
|
||||
], axis=1)
|
||||
for sid, df in enumerate(consensus_estimates_cases)}
|
||||
|
||||
loader_type = ConsensusEstimatesLoader
|
||||
|
||||
def setup(self, dates):
|
||||
cols = {
|
||||
PREVIOUS_RELEASE_DATE:
|
||||
self.get_expected_previous_event_dates(dates),
|
||||
NEXT_RELEASE_DATE: self.get_expected_next_event_dates(dates)
|
||||
}
|
||||
for field_name in field_name_to_expected_col:
|
||||
cols[field_name] = self.get_sids_to_frames(
|
||||
zip_with_floats, field_name_to_expected_col[field_name],
|
||||
self.prev_date_intervals
|
||||
if field_name.startswith("previous")
|
||||
else self.next_date_intervals,
|
||||
dates
|
||||
)
|
||||
return cols
|
||||
|
||||
|
||||
class BlazeConsensusEstimatesLoaderTestCase(ConsensusEstimatesLoaderTestCase):
|
||||
loader_type = BlazeConsensusEstimatesLoader
|
||||
|
||||
def pipeline_event_loader_args(self, dates):
|
||||
_, mapping = super(
|
||||
BlazeConsensusEstimatesLoaderTestCase,
|
||||
self,
|
||||
).pipeline_event_loader_args(dates)
|
||||
frames = []
|
||||
for sid, df in iteritems(mapping):
|
||||
frame = df.copy()
|
||||
frame[SID_FIELD_NAME] = sid
|
||||
frames.append(frame)
|
||||
return bz.data(pd.concat(frames).reset_index(drop=True)),
|
||||
|
||||
|
||||
class BlazeConsensusEstimatesLoaderNotInteractiveTestCase(
|
||||
BlazeConsensusEstimatesLoaderTestCase
|
||||
):
|
||||
"""Test case for passing a non-interactive symbol and a dict of resources.
|
||||
"""
|
||||
|
||||
def pipeline_event_loader_args(self, dates):
|
||||
(bound_expr,) = super(
|
||||
BlazeConsensusEstimatesLoaderNotInteractiveTestCase,
|
||||
self,
|
||||
).pipeline_event_loader_args(dates)
|
||||
return swap_resources_into_scope(bound_expr, {})
|
||||
@@ -45,7 +45,6 @@ from zipline.pipeline.loaders.dividends import (
|
||||
DividendsByPayDateLoader
|
||||
)
|
||||
from zipline.pipeline.loaders.utils import (
|
||||
get_values_for_date_ranges,
|
||||
zip_with_dates,
|
||||
zip_with_floats
|
||||
)
|
||||
@@ -159,33 +158,6 @@ next_amounts = [['NaN', 1, 15, 'NaN'],
|
||||
['NaN', 6, 23, 'NaN']]
|
||||
|
||||
|
||||
def get_vals_for_dates(zip_date_index_with_vals,
|
||||
vals,
|
||||
date_invervals,
|
||||
dates):
|
||||
return pd.DataFrame({
|
||||
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(WithPipelineEventDataLoader,
|
||||
ZiplineTestCase):
|
||||
"""
|
||||
@@ -232,14 +204,14 @@ class DividendsByAnnouncementDateTestCase(WithPipelineEventDataLoader,
|
||||
['NaT', '2014-01-04', '2014-01-14'],
|
||||
['NaT', '2014-01-04']]
|
||||
amounts = [['NaN', 1, 15], ['NaN', 7, 13], ['NaN', 3, 1], ['NaN', 23]]
|
||||
cols = {}
|
||||
cols[PREVIOUS_ANNOUNCEMENT] = get_vals_for_dates(
|
||||
zip_with_dates, announcement_dates, date_intervals, dates
|
||||
)
|
||||
|
||||
cols[PREVIOUS_AMOUNT] = get_vals_for_dates(
|
||||
zip_with_floats, amounts, date_intervals, dates
|
||||
)
|
||||
cols = {
|
||||
PREVIOUS_ANNOUNCEMENT: self.get_sids_to_frames(
|
||||
zip_with_dates, announcement_dates, date_intervals, dates
|
||||
),
|
||||
PREVIOUS_AMOUNT: self.get_sids_to_frames(
|
||||
zip_with_floats, amounts, date_intervals, dates
|
||||
),
|
||||
}
|
||||
|
||||
cols[
|
||||
DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT
|
||||
@@ -305,27 +277,26 @@ class DividendsByExDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase):
|
||||
loader_type = DividendsByExDateLoader
|
||||
|
||||
def setup(self, dates):
|
||||
cols = {}
|
||||
cols[NEXT_EX_DATE] = get_vals_for_dates(
|
||||
zip_with_dates, next_ex_and_pay_dates, next_date_intervals, dates,
|
||||
)
|
||||
|
||||
cols[PREVIOUS_EX_DATE] = get_vals_for_dates(
|
||||
zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals, dates
|
||||
)
|
||||
|
||||
cols[NEXT_AMOUNT] = get_vals_for_dates(
|
||||
zip_with_floats, next_amounts, next_date_intervals, dates
|
||||
)
|
||||
|
||||
cols[PREVIOUS_AMOUNT] = get_vals_for_dates(
|
||||
zip_with_floats, prev_amounts, prev_date_intervals, dates
|
||||
)
|
||||
cols = {
|
||||
NEXT_EX_DATE: self.get_sids_to_frames(
|
||||
zip_with_dates, next_ex_and_pay_dates, next_date_intervals,
|
||||
dates,
|
||||
),
|
||||
PREVIOUS_EX_DATE: self.get_sids_to_frames(
|
||||
zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals,
|
||||
dates
|
||||
),
|
||||
NEXT_AMOUNT: self.get_sids_to_frames(
|
||||
zip_with_floats, next_amounts, next_date_intervals, dates
|
||||
),
|
||||
PREVIOUS_AMOUNT: self.get_sids_to_frames(
|
||||
zip_with_floats, prev_amounts, prev_date_intervals, dates
|
||||
)
|
||||
}
|
||||
|
||||
cols[DAYS_TO_NEXT_EX_DATE] = self._compute_busday_offsets(
|
||||
cols[NEXT_EX_DATE]
|
||||
)
|
||||
|
||||
cols[DAYS_SINCE_PREV_EX_DATE] = self._compute_busday_offsets(
|
||||
cols[PREVIOUS_EX_DATE]
|
||||
)
|
||||
@@ -386,22 +357,22 @@ class DividendsByPayDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase):
|
||||
loader_type = DividendsByPayDateLoader
|
||||
|
||||
def setup(self, dates):
|
||||
cols = {}
|
||||
cols[NEXT_PAY_DATE] = get_vals_for_dates(
|
||||
zip_with_dates, next_ex_and_pay_dates, next_date_intervals, dates
|
||||
)
|
||||
cols[PREVIOUS_PAY_DATE] = get_vals_for_dates(
|
||||
zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals, dates
|
||||
)
|
||||
|
||||
cols[NEXT_AMOUNT] = get_vals_for_dates(
|
||||
zip_with_floats, next_amounts, next_date_intervals, dates
|
||||
)
|
||||
|
||||
cols[PREVIOUS_AMOUNT] = get_vals_for_dates(
|
||||
zip_with_floats, prev_amounts, prev_date_intervals, dates
|
||||
)
|
||||
return cols
|
||||
return {
|
||||
NEXT_PAY_DATE: self.get_sids_to_frames(
|
||||
zip_with_dates, next_ex_and_pay_dates, next_date_intervals,
|
||||
dates
|
||||
),
|
||||
PREVIOUS_PAY_DATE: self.get_sids_to_frames(
|
||||
zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals,
|
||||
dates
|
||||
),
|
||||
NEXT_AMOUNT: self.get_sids_to_frames(
|
||||
zip_with_floats, next_amounts, next_date_intervals, dates
|
||||
),
|
||||
PREVIOUS_AMOUNT: self.get_sids_to_frames(
|
||||
zip_with_floats, prev_amounts, prev_date_intervals, dates
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
class BlazeDividendsByPayDateLoaderTestCase(DividendsByPayDateTestCase):
|
||||
|
||||
+16
-145
@@ -22,97 +22,13 @@ from zipline.pipeline.factors.events import (
|
||||
)
|
||||
from zipline.pipeline.loaders.earnings import EarningsCalendarLoader
|
||||
from zipline.pipeline.loaders.blaze import BlazeEarningsCalendarLoader
|
||||
from zipline.pipeline.loaders.utils import (
|
||||
get_values_for_date_ranges,
|
||||
zip_with_dates
|
||||
)
|
||||
from zipline.testing.fixtures import (
|
||||
WithPipelineEventDataLoader,
|
||||
ZiplineTestCase
|
||||
ZiplineTestCase,
|
||||
WithNextAndPreviousEventDataLoader
|
||||
)
|
||||
|
||||
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]'
|
||||
),
|
||||
]
|
||||
|
||||
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]]
|
||||
]
|
||||
|
||||
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(WithPipelineEventDataLoader,
|
||||
class EarningsCalendarLoaderTestCase(WithNextAndPreviousEventDataLoader,
|
||||
ZiplineTestCase):
|
||||
"""
|
||||
Tests for loading the earnings announcement data.
|
||||
@@ -126,70 +42,25 @@ class EarningsCalendarLoaderTestCase(WithPipelineEventDataLoader,
|
||||
|
||||
@classmethod
|
||||
def get_dataset(cls):
|
||||
return {sid: df for sid, df in enumerate(earnings_cases)}
|
||||
return {sid: df.rename(
|
||||
columns={'other_date': ANNOUNCEMENT_FIELD_NAME}
|
||||
) for sid, df in enumerate(cls.base_cases)}
|
||||
|
||||
loader_type = EarningsCalendarLoader
|
||||
|
||||
def get_expected_next_event_dates(self, dates):
|
||||
return pd.DataFrame({
|
||||
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):
|
||||
return pd.DataFrame({
|
||||
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)
|
||||
|
||||
def setup(self, dates):
|
||||
_expected_next_announce = self.get_expected_next_event_dates(dates)
|
||||
|
||||
_expected_previous_announce = self.get_expected_previous_event_dates(
|
||||
dates
|
||||
cols = {
|
||||
PREVIOUS_ANNOUNCEMENT: self.get_expected_previous_event_dates(
|
||||
dates
|
||||
),
|
||||
NEXT_ANNOUNCEMENT: self.get_expected_next_event_dates(dates),
|
||||
}
|
||||
cols[DAYS_TO_NEXT] = self._compute_busday_offsets(
|
||||
cols[NEXT_ANNOUNCEMENT]
|
||||
)
|
||||
|
||||
_expected_next_busday_offsets = self._compute_busday_offsets(
|
||||
_expected_next_announce
|
||||
cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
|
||||
cols[PREVIOUS_ANNOUNCEMENT]
|
||||
)
|
||||
_expected_previous_busday_offsets = self._compute_busday_offsets(
|
||||
_expected_previous_announce
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user