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

MAINT: reorganize - move testing mixin to fixtures

Browse Source 
BUG: correctly create asset finder

MAINT: rename fixture

STY: fixes for flake8

STY: add space around assignment

MAINT: add var back to constructor

MAINT: remove unused import

MAINT: compare var with None directly

MAINT: fix merge errors
This commit is contained in:
Maya Tydykov
2016-03-23 15:41:22 -04:00
parent 6b3560ade8
commit e8185a1512
13 changed files with 205 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/base.py
+1 -112
View File
@@ -1,34 +1,25 @@
"""
Base class for Pipeline API unittests.
"""
import abc
from functools import wraps
from unittest import TestCase
from nose_parameterized import parameterized
import numpy as np
from numpy import arange, prod
import pandas as pd
from pandas import date_range, Int64Index, DataFrame
from pandas.util.testing import assert_series_equal
from six import iteritems
from zipline.pipeline import Pipeline, TermGraph
from zipline.pipeline import TermGraph
from zipline.pipeline.engine import SimplePipelineEngine
from zipline.pipeline.term import AssetExists
from zipline.testing import (
check_arrays,
ExplodingObject,
gen_calendars,
make_simple_equity_info,
tmp_asset_finder,
)
from zipline.utils.functional import dzip_exact
from zipline.utils.numpy_utils import (
NaTD,
make_datetime64D
)
from zipline.utils.pandas_utils import explode
from zipline.utils.tradingcalendar import trading_day
@@ -175,105 +166,3 @@ class BasePipelineTestCase(TestCase):
@with_default_shape
def ones_mask(self, shape):
return np.ones(shape, dtype=bool)
class EventLoaderCommonMixin(object):
@abc.abstractproperty
def get_sids(cls):
raise NotImplementedError('get_sids')
@abc.abstractproperty
def get_dataset(self):
raise NotImplementedError('get_dataset')
@abc.abstractproperty
def loader_type(self):
raise NotImplementedError('loader_type')
def loader_args(self, dates):
"""Construct the base object to pass to the loader.
Parameters
----------
dates : pd.DatetimeIndex
The dates we can serve.
Returns
-------
args : tuple[any]
The arguments to forward to the loader positionally.
"""
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.asset_finder)
@staticmethod
def _compute_busday_offsets(announcement_dates):
"""
Compute expected business day offsets from a DataFrame of announcement
dates.
"""
# Column-vector of dates on which factor `compute` will be called.
raw_call_dates = announcement_dates.index.values.astype(
'datetime64[D]'
)[:, None]
# 2D array of dates containining expected nexg announcement.
raw_announce_dates = (
announcement_dates.values.astype('datetime64[D]')
)
# Set NaTs to 0 temporarily because busday_count doesn't support NaT.
# We fill these entries with NaNs later.
whereNaT = raw_announce_dates == NaTD
raw_announce_dates[whereNaT] = make_datetime64D(0)
# The abs call here makes it so that we can use this function to
# compute offsets for both next and previous earnings (previous
# earnings offsets come back negative).
expected = abs(np.busday_count(
raw_call_dates,
raw_announce_dates
).astype(float))
expected[whereNaT] = np.nan
return pd.DataFrame(
data=expected,
columns=announcement_dates.columns,
index=announcement_dates.index,
)
@parameterized.expand(gen_calendars(
'2014-01-01',
'2014-01-31',
critical_dates=pd.to_datetime([
'2014-01-05',
'2014-01-10',
'2014-01-15',
'2014-01-20',
], utc=True),
))
def test_compute(self, dates):
engine = self.setup_engine(dates)
cols = self.setup(dates)
pipe = Pipeline(
columns=self.pipeline_columns
)
result = engine.run_pipeline(
pipe,
start_date=dates[0],
end_date=dates[-1],
)
for sid in self.get_sids():
for col_name in cols.keys():
assert_series_equal(result[col_name].xs(sid, level=1),
cols[col_name][sid],
check_names=False)
tests/pipeline/test_buyback_auth.py
+15 -14
View File
@@ -5,7 +5,6 @@ import blaze as bz
from blaze.compute.core import swap_resources_into_scope
import pandas as pd
from six import iteritems
from .base import EventLoaderCommonMixin
from zipline.pipeline.common import(
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
@@ -39,7 +38,9 @@ from zipline.pipeline.loaders.utils import (
zip_with_floats,
zip_with_dates
)
from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase
from zipline.testing.fixtures import (
WithPipelineEventDataLoader, ZiplineTestCase
)
date_intervals = [[None, '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', None]]
@@ -74,8 +75,8 @@ def get_expected_previous_values(zip_date_index_with_vals,
}, index=dates)
class CashBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase,
EventLoaderCommonMixin):
class CashBuybackAuthLoaderTestCase(WithPipelineEventDataLoader,
ZiplineTestCase):
"""
Test for cash buyback authorizations dataset.
"""
@@ -118,8 +119,8 @@ class CashBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase,
return cols
class ShareBuybackAuthLoaderTestCase(WithAssetFinder, ZiplineTestCase,
EventLoaderCommonMixin):
class ShareBuybackAuthLoaderTestCase(WithPipelineEventDataLoader,
ZiplineTestCase):
"""
Test for share buyback authorizations dataset.
"""
@@ -167,11 +168,11 @@ class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase):
"""
loader_type = BlazeCashBuybackAuthorizationsLoader
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeCashBuybackAuthLoaderTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return (bz.data(pd.concat(
pd.DataFrame({
BUYBACK_ANNOUNCEMENT_FIELD_NAME:
@@ -191,11 +192,11 @@ class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase):
"""
loader_type = BlazeShareBuybackAuthorizationsLoader
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeShareBuybackAuthLoaderTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return (bz.data(pd.concat(
pd.DataFrame({
BUYBACK_ANNOUNCEMENT_FIELD_NAME:
@@ -214,11 +215,11 @@ class BlazeShareBuybackAuthLoaderNotInteractiveTestCase(
BlazeShareBuybackAuthLoaderTestCase):
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeShareBuybackAuthLoaderNotInteractiveTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
@@ -226,9 +227,9 @@ class BlazeCashBuybackAuthLoaderNotInteractiveTestCase(
BlazeCashBuybackAuthLoaderTestCase):
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeCashBuybackAuthLoaderNotInteractiveTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
tests/pipeline/test_dividends.py
+24 -36
View File
@@ -5,7 +5,6 @@ import blaze as bz
from blaze.compute.core import swap_resources_into_scope
import pandas as pd
from six import iteritems
from tests.pipeline.base import EventLoaderCommonMixin
from zipline.pipeline.common import (
ANNOUNCEMENT_FIELD_NAME,
@@ -50,7 +49,10 @@ from zipline.pipeline.loaders.utils import (
zip_with_dates,
zip_with_floats
)
from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase
from zipline.testing.fixtures import (
WithPipelineEventDataLoader,
ZiplineTestCase
)
dividends_cases = [
# K1--K2--A1--A2.
@@ -184,8 +186,8 @@ def get_vals_for_dates(zip_date_index_with_vals,
}, index=dates)
class DividendsByAnnouncementDateTestCase(WithAssetFinder, ZiplineTestCase,
EventLoaderCommonMixin):
class DividendsByAnnouncementDateTestCase(WithPipelineEventDataLoader,
ZiplineTestCase):
"""
Tests for loading the dividends by announcement date data.
"""
@@ -197,10 +199,6 @@ class DividendsByAnnouncementDateTestCase(WithAssetFinder, ZiplineTestCase,
BusinessDaysSinceDividendAnnouncement(),
}
@classmethod
def get_sids(cls):
return range(0, 5)
@classmethod
def get_dataset(cls):
return {sid:
@@ -254,11 +252,11 @@ class BlazeDividendsByAnnouncementDateTestCase(
):
loader_type = BlazeDividendsByAnnouncementDateLoader
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeDividendsByAnnouncementDateTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return (bz.Data(pd.concat(
pd.DataFrame({
ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME],
@@ -275,32 +273,27 @@ class BlazeDividendsByAnnouncementDateNotInteractiveTestCase(
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeDividendsByAnnouncementDateNotInteractiveTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
class DividendsByExDateTestCase(WithAssetFinder, ZiplineTestCase,
EventLoaderCommonMixin):
class DividendsByExDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase):
"""
Tests for loading the dividends by ex date data.
"""
pipeline_columns = {
NEXT_EX_DATE: DividendsByExDate.next_ex_date.latest,
PREVIOUS_EX_DATE: DividendsByExDate.previous_ex_date.latest,
NEXT_EX_DATE: DividendsByExDate.next_date.latest,
PREVIOUS_EX_DATE: DividendsByExDate.previous_date.latest,
NEXT_AMOUNT: DividendsByExDate.next_amount.latest,
PREVIOUS_AMOUNT: DividendsByExDate.previous_amount.latest,
DAYS_TO_NEXT_EX_DATE: BusinessDaysUntilNextExDate(),
DAYS_SINCE_PREV_EX_DATE: BusinessDaysSincePreviousExDate()
}
@classmethod
def get_sids(cls):
return range(0, 5)
@classmethod
def get_dataset(cls):
return {sid:
@@ -342,11 +335,11 @@ class DividendsByExDateTestCase(WithAssetFinder, ZiplineTestCase,
class BlazeDividendsByExDateLoaderTestCase(DividendsByExDateTestCase):
loader_type = BlazeDividendsByExDateLoader
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeDividendsByExDateLoaderTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return (bz.Data(pd.concat(
pd.DataFrame({
EX_DATE_FIELD_NAME: df[EX_DATE_FIELD_NAME],
@@ -363,30 +356,25 @@ class BlazeDividendsByExDateLoaderNotInteractiveTestCase(
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeDividendsByExDateLoaderNotInteractiveTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
class DividendsByPayDateTestCase(WithAssetFinder, ZiplineTestCase,
EventLoaderCommonMixin):
class DividendsByPayDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase):
"""
Tests for loading the dividends by pay date data.
"""
pipeline_columns = {
NEXT_PAY_DATE: DividendsByPayDate.next_pay_date.latest,
PREVIOUS_PAY_DATE: DividendsByPayDate.previous_pay_date.latest,
NEXT_PAY_DATE: DividendsByPayDate.next_date.latest,
PREVIOUS_PAY_DATE: DividendsByPayDate.previous_date.latest,
NEXT_AMOUNT: DividendsByPayDate.next_amount.latest,
PREVIOUS_AMOUNT: DividendsByPayDate.previous_amount.latest,
}
@classmethod
def get_sids(cls):
return range(0, 5)
@classmethod
def get_dataset(cls):
return {sid:
@@ -419,11 +407,11 @@ class DividendsByPayDateTestCase(WithAssetFinder, ZiplineTestCase,
class BlazeDividendsByPayDateLoaderTestCase(DividendsByPayDateTestCase):
loader_type = BlazeDividendsByPayDateLoader
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeDividendsByPayDateLoaderTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return (bz.Data(pd.concat(
pd.DataFrame({
PAY_DATE_FIELD_NAME: df[PAY_DATE_FIELD_NAME],
@@ -440,9 +428,9 @@ class BlazeDividendsByPayDateLoaderNotInteractiveTestCase(
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeDividendsByPayDateLoaderNotInteractiveTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
tests/pipeline/test_earnings.py
+10 -12
View File
@@ -5,7 +5,6 @@ import blaze as bz
from blaze.compute.core import swap_resources_into_scope
import pandas as pd
from six import iteritems
from .base import EventLoaderCommonMixin
from zipline.pipeline.common import (
ANNOUNCEMENT_FIELD_NAME,
@@ -28,7 +27,10 @@ from zipline.pipeline.loaders.utils import (
zip_with_dates
)
from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase
from zipline.testing.fixtures import (
WithPipelineEventDataLoader,
ZiplineTestCase
)
earnings_cases = [
# K1--K2--A1--A2.
@@ -111,8 +113,8 @@ prev_dates = [
]
class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase,
EventLoaderCommonMixin):
class EarningsCalendarLoaderTestCase(WithPipelineEventDataLoader,
ZiplineTestCase):
"""
Tests for loading the earnings announcement data.
"""
@@ -123,10 +125,6 @@ class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase,
DAYS_TO_NEXT: BusinessDaysUntilNextEarnings(),
}
@classmethod
def get_sids(cls):
return range(5)
@classmethod
def get_dataset(cls):
return {sid: df for sid, df in enumerate(earnings_cases)}
@@ -199,11 +197,11 @@ class EarningsCalendarLoaderTestCase(WithAssetFinder, ZiplineTestCase,
class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase):
loader_type = BlazeEarningsCalendarLoader
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeEarningsCalendarLoaderTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return (bz.data(pd.concat(
pd.DataFrame({
ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME],
@@ -219,9 +217,9 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase(
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def loader_args(self, dates):
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeEarningsCalendarLoaderNotInteractiveTestCase,
self,
).loader_args(dates)
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
zipline/pipeline/data/dividends.py
+4 -4
View File
@@ -7,15 +7,15 @@ from .dataset import Column, DataSet
class DividendsByExDate(DataSet):
next_ex_date = Column(datetime64ns_dtype)
previous_ex_date = Column(datetime64ns_dtype)
next_date = Column(datetime64ns_dtype)
previous_date = Column(datetime64ns_dtype)
next_amount = Column(float64_dtype)
previous_amount = Column(float64_dtype)
class DividendsByPayDate(DataSet):
next_pay_date = Column(datetime64ns_dtype)
previous_pay_date = Column(datetime64ns_dtype)
next_date = Column(datetime64ns_dtype)
previous_date = Column(datetime64ns_dtype)
next_amount = Column(float64_dtype)
previous_amount = Column(float64_dtype)
zipline/pipeline/factors/events.py
+2 -2
View File
@@ -189,7 +189,7 @@ class BusinessDaysUntilNextExDate(
--------
zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement
"""
inputs = [DividendsByExDate.next_ex_date]
inputs = [DividendsByExDate.next_date]
class BusinessDaysSincePreviousExDate(
@@ -204,4 +204,4 @@ class BusinessDaysSincePreviousExDate(
--------
zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement
"""
inputs = [DividendsByExDate.previous_ex_date]
inputs = [DividendsByExDate.previous_date]
zipline/pipeline/loaders/blaze/buyback_auth.py
+2 -36
View File
@@ -68,24 +68,7 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader):
})
concrete_loader = CashBuybackAuthorizationsLoader
def __init__(self,
expr,
resources=None,
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=CashBuybackAuthorizations,
**kwargs):
super(
BlazeCashBuybackAuthorizationsLoader, self
).__init__(expr,
resources=resources,
odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz,
dataset=dataset,
**kwargs)
default_dataset = CashBuybackAuthorizations
class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader):
@@ -140,21 +123,4 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader):
})
concrete_loader = ShareBuybackAuthorizationsLoader
def __init__(self,
expr,
resources=None,
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=ShareBuybackAuthorizations,
**kwargs):
super(
BlazeShareBuybackAuthorizationsLoader, self
).__init__(expr,
resources=resources,
odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz,
dataset=dataset,
**kwargs)
default_dataset = ShareBuybackAuthorizations
zipline/pipeline/loaders/blaze/dividends.py
+3 -45
View File
@@ -72,21 +72,7 @@ class BlazeDividendsByAnnouncementDateLoader(BlazeEventsLoader):
})
concrete_loader = DividendsByAnnouncementDateLoader
def __init__(self,
expr,
resources=None,
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=DividendsByAnnouncementDate,
**kwargs):
super(
BlazeDividendsByAnnouncementDateLoader, self
).__init__(expr, dataset=dataset,
resources=resources, odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz, **kwargs)
default_dataset = DividendsByAnnouncementDate
class BlazeDividendsByExDateLoader(BlazeEventsLoader):
@@ -142,21 +128,7 @@ class BlazeDividendsByExDateLoader(BlazeEventsLoader):
})
concrete_loader = DividendsByExDateLoader
def __init__(self,
expr,
resources=None,
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=DividendsByExDate,
**kwargs):
super(
BlazeDividendsByExDateLoader, self
).__init__(expr, dataset=dataset,
resources=resources, odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz, **kwargs)
default_dataset = DividendsByExDate
class BlazeDividendsByPayDateLoader(BlazeEventsLoader):
@@ -212,18 +184,4 @@ class BlazeDividendsByPayDateLoader(BlazeEventsLoader):
})
concrete_loader = DividendsByPayDateLoader
def __init__(self,
expr,
resources=None,
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=DividendsByPayDate,
**kwargs):
super(
BlazeDividendsByPayDateLoader, self
).__init__(expr, dataset=dataset,
resources=resources, odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz, **kwargs)
default_dataset = DividendsByPayDate
zipline/pipeline/loaders/blaze/earnings.py
+1 -15
View File
@@ -58,18 +58,4 @@ class BlazeEarningsCalendarLoader(BlazeEventsLoader):
})
concrete_loader = EarningsCalendarLoader
def __init__(self,
expr,
resources=None,
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=EarningsCalendar,
**kwargs):
super(
BlazeEarningsCalendarLoader, self
).__init__(expr, dataset=dataset,
resources=resources, odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz, **kwargs)
default_dataset = EarningsCalendar
zipline/pipeline/loaders/blaze/events.py
+5 -1
View File
@@ -56,6 +56,7 @@ class BlazeEventsLoader(PipelineLoader):
If the '{TS_FIELD_NAME}' field is not included it is assumed that we
start the backtest with knowledge of all announcements.
"""
default_dataset = None
@preprocess(data_query_tz=optionally(ensure_timezone))
def __init__(self,
@@ -64,7 +65,10 @@ class BlazeEventsLoader(PipelineLoader):
odo_kwargs=None,
data_query_time=None,
data_query_tz=None,
dataset=None):
dataset=default_dataset):
if dataset is None:
dataset = self.default_dataset
dshape = expr.dshape
if not istabular(dshape):
zipline/pipeline/loaders/dividends.py
+8 -8
View File
@@ -52,14 +52,14 @@ class DividendsByPayDateLoader(EventsLoader):
)
@lazyval
def next_pay_date_loader(self):
return self._next_event_date_loader(self.dataset.next_pay_date,
def next_date_loader(self):
return self._next_event_date_loader(self.dataset.next_date,
PAY_DATE_FIELD_NAME)
@lazyval
def previous_pay_date_loader(self):
def previous_date_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_pay_date,
self.dataset.previous_date,
PAY_DATE_FIELD_NAME
)
@@ -90,14 +90,14 @@ class DividendsByExDateLoader(EventsLoader):
)
@lazyval
def next_ex_date_loader(self):
return self._next_event_date_loader(self.dataset.next_ex_date,
def next_date_loader(self):
return self._next_event_date_loader(self.dataset.next_date,
EX_DATE_FIELD_NAME)
@lazyval
def previous_ex_date_loader(self):
def previous_date_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_ex_date,
self.dataset.previous_date,
EX_DATE_FIELD_NAME
)
zipline/pipeline/loaders/utils.py
+3 -2
View File
@@ -60,8 +60,9 @@ def next_event_frame(events_by_sid,
# Iterate over the raw Series values, since we're comparing against
# numpy arrays anyway.
iterkv = zip(event_dates.index.values, event_dates.values, values)
for knowledge_date, event_date, value in iterkv:
iter_date_vals = zip(event_dates.index.values, event_dates.values,
values)
for knowledge_date, event_date, value in iter_date_vals:
date_mask = (
(knowledge_date <= raw_dates) &
(raw_dates <= event_date)
zipline/testing/fixtures.py
+127 -9
View File
@@ -2,13 +2,19 @@ from unittest import TestCase
from contextlib2 import ExitStack
from logbook import NullHandler
from nose_parameterized import parameterized
import numpy as np
import pandas as pd
from pandas.util.testing import assert_series_equal
from six import with_metaclass
from .core import tmp_asset_finder, make_simple_equity_info
from .core import tmp_asset_finder, make_simple_equity_info, gen_calendars
from ..finance.trading import TradingEnvironment
from ..utils import tradingcalendar, factory
from ..utils.final import FinalMeta, final
from zipline.pipeline import Pipeline, SimplePipelineEngine
from zipline.utils.numpy_utils import make_datetime64D
from zipline.utils.numpy_utils import NaTD
class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
@@ -177,14 +183,7 @@ class WithAssetFinder(object):
def _make_info(cls):
return None
@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_equities_info = _make_info
make_futures_info = _make_info
make_exchanges_info = _make_info
make_root_symbols_info = _make_info
@@ -299,3 +298,122 @@ class WithNYSETradingDays(object):
start_loc = end_loc - cls.TRADING_DAY_COUNT
cls.trading_days = all_days[start_loc:end_loc + 1]
class WithPipelineEventDataLoader(WithAssetFinder):
"""
ZiplineTestCase mixin providing common test methods/behaviors for event
data loaders.
`get_sids` must return the sids being tested.
`get_dataset` must return {sid -> pd.DataFrame}
`loader_type` must return the loader class to use for loading the dataset
`make_asset_finder` returns a default asset finder which can be overridden.
"""
@classmethod
def get_sids(cls):
return range(0, 5)
@classmethod
def get_dataset(cls):
return {sid: pd.DataFrame() for sid in cls.get_sids()}
@classmethod
def loader_type(self):
return None
@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'),
)
def pipeline_event_loader_args(self, dates):
"""Construct the base object to pass to the loader.
Parameters
----------
dates : pd.DatetimeIndex
The dates we can serve.
Returns
-------
args : tuple[any]
The arguments to forward to the loader positionally.
"""
return dates, self.get_dataset()
def pipeline_event_setup_engine(self, dates):
"""
Make a Pipeline Enigne object based on the given dates.
"""
loader = self.loader_type(*self.pipeline_event_loader_args(dates))
return SimplePipelineEngine(lambda _: loader, dates, self.asset_finder)
@staticmethod
def _compute_busday_offsets(announcement_dates):
"""
Compute expected business day offsets from a DataFrame of announcement
dates.
"""
# Column-vector of dates on which factor `compute` will be called.
raw_call_dates = announcement_dates.index.values.astype(
'datetime64[D]'
)[:, None]
# 2D array of dates containining expected nexg announcement.
raw_announce_dates = (
announcement_dates.values.astype('datetime64[D]')
)
# Set NaTs to 0 temporarily because busday_count doesn't support NaT.
# We fill these entries with NaNs later.
whereNaT = raw_announce_dates == NaTD
raw_announce_dates[whereNaT] = make_datetime64D(0)
# The abs call here makes it so that we can use this function to
# compute offsets for both next and previous earnings (previous
# earnings offsets come back negative).
expected = abs(np.busday_count(
raw_call_dates,
raw_announce_dates
).astype(float))
expected[whereNaT] = np.nan
return pd.DataFrame(
data=expected,
columns=announcement_dates.columns,
index=announcement_dates.index,
)
@parameterized.expand(gen_calendars(
'2014-01-01',
'2014-01-31',
critical_dates=pd.to_datetime([
'2014-01-05',
'2014-01-10',
'2014-01-15',
'2014-01-20',
], utc=True),
))
def test_compute(self, dates):
engine = self.pipeline_event_setup_engine(dates)
cols = self.setup(dates)
pipe = Pipeline(
columns=self.pipeline_columns
)
result = engine.run_pipeline(
pipe,
start_date=dates[0],
end_date=dates[-1],
)
for sid in self.get_sids():
for col_name in cols.keys():
assert_series_equal(result[col_name].xs(sid, level=1),
cols[col_name][sid],
check_names=False)