mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-03 02:28:28 +08:00
0191d9d903
TST: test both next and prev event frame loading and use EventsLoader. BUG: remove extra arg MAINT: call list on zip for compatibility with python 3
1179 lines
40 KiB
Python
1179 lines
40 KiB
Python
from abc import ABCMeta, abstractproperty
|
|
import sqlite3
|
|
from unittest import TestCase
|
|
|
|
from contextlib2 import ExitStack
|
|
from logbook import NullHandler, Logger
|
|
from nose_parameterized import parameterized
|
|
from pandas.util.testing import assert_series_equal
|
|
from six import with_metaclass, iteritems
|
|
from toolz import flip
|
|
import numpy as np
|
|
import pandas as pd
|
|
import responses
|
|
|
|
|
|
from .core import (
|
|
create_daily_bar_data,
|
|
create_minute_bar_data,
|
|
tmp_dir,
|
|
)
|
|
from ..data.data_portal import DataPortal
|
|
from ..data.us_equity_pricing import (
|
|
SQLiteAdjustmentReader,
|
|
SQLiteAdjustmentWriter,
|
|
)
|
|
from ..data.us_equity_pricing import (
|
|
BcolzDailyBarReader,
|
|
BcolzDailyBarWriter,
|
|
)
|
|
from ..data.minute_bars import (
|
|
BcolzMinuteBarReader,
|
|
BcolzMinuteBarWriter,
|
|
US_EQUITIES_MINUTES_PER_DAY
|
|
)
|
|
|
|
from ..finance.trading import TradingEnvironment
|
|
from ..utils import tradingcalendar, factory
|
|
from ..utils.classproperty import classproperty
|
|
from ..utils.final import FinalMeta, final
|
|
from ..utils.metautils import compose_types
|
|
from .core import tmp_asset_finder, make_simple_equity_info, gen_calendars
|
|
from zipline.pipeline import Pipeline, SimplePipelineEngine
|
|
from zipline.utils.numpy_utils import make_datetime64D
|
|
from zipline.utils.numpy_utils import NaTD
|
|
from zipline.pipeline.common import TS_FIELD_NAME
|
|
from zipline.pipeline.loaders.utils import (
|
|
get_values_for_date_ranges,
|
|
zip_with_dates
|
|
)
|
|
|
|
|
|
class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
|
|
"""
|
|
Shared extensions to core unittest.TestCase.
|
|
|
|
Overrides the default unittest setUp/tearDown functions with versions that
|
|
use ExitStack to correctly clean up resources, even in the face of
|
|
exceptions that occur during setUp/setUpClass.
|
|
|
|
Subclasses **should not override setUp or setUpClass**!
|
|
|
|
Instead, they should implement `init_instance_fixtures` for per-test-method
|
|
resources, and `init_class_fixtures` for per-class resources.
|
|
|
|
Resources that need to be cleaned up should be registered using
|
|
either `enter_{class,instance}_context` or `add_{class,instance}_callback}.
|
|
"""
|
|
_in_setup = False
|
|
|
|
@final
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
# Hold a set of all the "static" attributes on the class. These are
|
|
# things that are not populated after the class was created like
|
|
# methods or other class level attributes.
|
|
cls._static_class_attributes = set(vars(cls))
|
|
cls._class_teardown_stack = ExitStack()
|
|
try:
|
|
cls._base_init_fixtures_was_called = False
|
|
cls.init_class_fixtures()
|
|
assert cls._base_init_fixtures_was_called, (
|
|
"ZiplineTestCase.init_class_fixtures() was not called.\n"
|
|
"This probably means that you overrode init_class_fixtures"
|
|
" without calling super()."
|
|
)
|
|
except:
|
|
cls.tearDownClass()
|
|
raise
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
"""
|
|
Override and implement this classmethod to register resources that
|
|
should be created and/or torn down on a per-class basis.
|
|
|
|
Subclass implementations of this should always invoke this with super()
|
|
to ensure that fixture mixins work properly.
|
|
"""
|
|
if cls._in_setup:
|
|
raise ValueError(
|
|
'Called init_class_fixtures from init_instance_fixtures.'
|
|
'Did you write super(..., self).init_class_fixtures() instead'
|
|
' of super(..., self).init_instance_fixtures()?',
|
|
)
|
|
cls._base_init_fixtures_was_called = True
|
|
|
|
@final
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._class_teardown_stack.close()
|
|
for name in set(vars(cls)) - cls._static_class_attributes:
|
|
# Remove all of the attributes that were added after the class was
|
|
# constructed. This cleans up any large test data that is class
|
|
# scoped while still allowing subclasses to access class level
|
|
# attributes.
|
|
delattr(cls, name)
|
|
|
|
@final
|
|
@classmethod
|
|
def enter_class_context(cls, context_manager):
|
|
"""
|
|
Enter a context manager to be exited during the tearDownClass
|
|
"""
|
|
if cls._in_setup:
|
|
raise ValueError(
|
|
'Attempted to enter a class context in init_instance_fixtures.'
|
|
'\nDid you mean to call enter_instance_context?',
|
|
)
|
|
return cls._class_teardown_stack.enter_context(context_manager)
|
|
|
|
@final
|
|
@classmethod
|
|
def add_class_callback(cls, callback):
|
|
"""
|
|
Register a callback to be executed during tearDownClass.
|
|
|
|
Parameters
|
|
----------
|
|
callback : callable
|
|
The callback to invoke at the end of the test suite.
|
|
"""
|
|
if cls._in_setup:
|
|
raise ValueError(
|
|
'Attempted to add a class callback in init_instance_fixtures.'
|
|
'\nDid you mean to call add_instance_callback?',
|
|
)
|
|
return cls._class_teardown_stack.callback(callback)
|
|
|
|
@final
|
|
def setUp(self):
|
|
type(self)._in_setup = True
|
|
self._pre_setup_attrs = set(vars(self))
|
|
self._instance_teardown_stack = ExitStack()
|
|
try:
|
|
self._init_instance_fixtures_was_called = False
|
|
self.init_instance_fixtures()
|
|
assert self._init_instance_fixtures_was_called, (
|
|
"ZiplineTestCase.init_instance_fixtures() was not"
|
|
" called.\n"
|
|
"This probably means that you overrode"
|
|
" init_instance_fixtures without calling super()."
|
|
)
|
|
except:
|
|
self.tearDown()
|
|
raise
|
|
finally:
|
|
type(self)._in_setup = False
|
|
|
|
def init_instance_fixtures(self):
|
|
self._init_instance_fixtures_was_called = True
|
|
|
|
@final
|
|
def tearDown(self):
|
|
self._instance_teardown_stack.close()
|
|
for attr in set(vars(self)) - self._pre_setup_attrs:
|
|
delattr(self, attr)
|
|
|
|
@final
|
|
def enter_instance_context(self, context_manager):
|
|
"""
|
|
Enter a context manager that should be exited during tearDown.
|
|
"""
|
|
return self._instance_teardown_stack.enter_context(context_manager)
|
|
|
|
@final
|
|
def add_instance_callback(self, callback):
|
|
"""
|
|
Register a callback to be executed during tearDown.
|
|
|
|
Parameters
|
|
----------
|
|
callback : callable
|
|
The callback to invoke at the end of each test.
|
|
"""
|
|
return self._instance_teardown_stack.callback(callback)
|
|
|
|
|
|
def alias(attr_name):
|
|
"""Make a fixture attribute an alias of another fixture's attribute by
|
|
default.
|
|
|
|
Parameters
|
|
----------
|
|
attr_name : str
|
|
The name of the attribute to alias.
|
|
|
|
Returns
|
|
-------
|
|
p : classproperty
|
|
A class property that does the property aliasing.
|
|
|
|
Examples
|
|
--------
|
|
>>> class C(object):
|
|
... attr = 1
|
|
...
|
|
>>> class D(object):
|
|
... attr_alias = alias('attr')
|
|
...
|
|
>>> D.attr
|
|
1
|
|
>>> D.attr_alias
|
|
1
|
|
>>> class E(D):
|
|
... attr_alias = 2
|
|
...
|
|
>>> E.attr
|
|
1
|
|
>>> E.attr_alias
|
|
2
|
|
"""
|
|
return classproperty(flip(getattr, attr_name))
|
|
|
|
|
|
class WithDefaultDateBounds(object):
|
|
"""
|
|
ZiplineTestCase mixin which makes it possible to synchronize date bounds
|
|
across fixtures.
|
|
|
|
Attributes
|
|
----------
|
|
START_DATE : datetime
|
|
END_DATE : datetime
|
|
The date bounds to be used for fixtures that want to have consistent
|
|
dates.
|
|
"""
|
|
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
|
|
END_DATE = pd.Timestamp('2006-12-29', tz='utc')
|
|
|
|
|
|
class WithLogger(object):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.log_handler as an instance-level
|
|
fixture.
|
|
|
|
After init_instance_fixtures has been called `self.log_handler` will be a
|
|
new ``logbook.NullHandler``.
|
|
|
|
Methods
|
|
-------
|
|
make_log_handler() -> logbook.LogHandler
|
|
A class method which constructs the new log handler object. By default
|
|
this will construct a ``NullHandler``.
|
|
"""
|
|
make_log_handler = NullHandler
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithLogger, cls).init_class_fixtures()
|
|
cls.log = Logger()
|
|
cls.log_handler = cls.enter_class_context(
|
|
cls.make_log_handler().applicationbound(),
|
|
)
|
|
|
|
|
|
class WithAssetFinder(WithDefaultDateBounds):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.asset_finder as a class-level fixture.
|
|
|
|
After init_class_fixtures has been called, `cls.asset_finder` is populated
|
|
with an AssetFinder.
|
|
|
|
Attributes
|
|
----------
|
|
ASSET_FINDER_EQUITY_SIDS : iterable[int]
|
|
The default sids to construct equity data for.
|
|
ASSET_FINDER_EQUITY_SYMBOLS : iterable[str]
|
|
The default symbols to use for the equities.
|
|
ASSET_FINDER_EQUITY_START_DATE : datetime
|
|
The default start date to create equity data for. This defaults to
|
|
``START_DATE``.
|
|
ASSET_FINDER_EQUITY_END_DATE : datetime
|
|
The default end date to create equity data for. This defaults to
|
|
``END_DATE``.
|
|
|
|
Methods
|
|
-------
|
|
make_equity_info() -> pd.DataFrame
|
|
A class method which constructs the dataframe of equity info to write
|
|
to the class's asset db. By default this is empty.
|
|
make_futures_info() -> pd.DataFrame
|
|
A class method which constructs the dataframe of futures contract info
|
|
to write to the class's asset db. By default this is empty.
|
|
make_exchanges_info() -> pd.DataFrame
|
|
A class method which constructs the dataframe of exchange information
|
|
to write to the class's assets db. By default this is empty.
|
|
make_root_symbols_info() -> pd.DataFrame
|
|
A class method which constructs the dataframe of root symbols
|
|
information to write to the class's assets db. By default this is
|
|
empty.
|
|
make_asset_finder() -> pd.DataFrame
|
|
A class method which constructs the actual asset finder object to use
|
|
for the class. If this method is overridden then the ``make_*_info``
|
|
methods may not be respected.
|
|
|
|
See Also
|
|
--------
|
|
zipline.testing.make_simple_equity_info
|
|
zipline.testing.make_jagged_equity_info
|
|
zipline.testing.make_rotating_equity_info
|
|
zipline.testing.make_future_info
|
|
zipline.testing.make_commodity_future_info
|
|
"""
|
|
ASSET_FINDER_EQUITY_SIDS = ord('A'), ord('B'), ord('C')
|
|
ASSET_FINDER_EQUITY_SYMBOLS = None
|
|
ASSET_FINDER_EQUITY_START_DATE = alias('START_DATE')
|
|
ASSET_FINDER_EQUITY_END_DATE = alias('END_DATE')
|
|
|
|
@classmethod
|
|
def _make_info(cls):
|
|
return None
|
|
|
|
make_futures_info = _make_info
|
|
make_exchanges_info = _make_info
|
|
make_root_symbols_info = _make_info
|
|
|
|
del _make_info
|
|
|
|
@classmethod
|
|
def make_equity_info(cls):
|
|
return make_simple_equity_info(
|
|
cls.ASSET_FINDER_EQUITY_SIDS,
|
|
cls.ASSET_FINDER_EQUITY_START_DATE,
|
|
cls.ASSET_FINDER_EQUITY_END_DATE,
|
|
cls.ASSET_FINDER_EQUITY_SYMBOLS,
|
|
)
|
|
|
|
@classmethod
|
|
def make_asset_finder(cls):
|
|
return cls.enter_class_context(tmp_asset_finder(
|
|
equities=cls.make_equity_info(),
|
|
futures=cls.make_futures_info(),
|
|
exchanges=cls.make_exchanges_info(),
|
|
root_symbols=cls.make_root_symbols_info(),
|
|
))
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithAssetFinder, cls).init_class_fixtures()
|
|
cls.asset_finder = cls.make_asset_finder()
|
|
|
|
|
|
class WithTradingEnvironment(WithAssetFinder):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.env as a class-level fixture.
|
|
|
|
After ``init_class_fixtures`` has been called, `cls.env` is populated
|
|
with a trading environment whose `asset_finder` is the result of
|
|
`cls.make_asset_finder`.
|
|
|
|
Attributes
|
|
----------
|
|
TRADING_ENV_MIN_DATE : datetime
|
|
The min_date to forward to the constructed TradingEnvironment.
|
|
TRADING_ENV_MAX_DATE : datetime
|
|
The max date to forward to the constructed TradingEnvironment.
|
|
TRADING_ENV_TRADING_CALENDAR : pd.DatetimeIndex
|
|
The trading calendar to use for the class's TradingEnvironment.
|
|
|
|
Methods
|
|
-------
|
|
make_load_function() -> callable
|
|
A class method that returns the ``load`` argument to pass to the
|
|
constructor of ``TradingEnvironment`` for this class.
|
|
The signature for the callable returned is:
|
|
``(datetime, pd.DatetimeIndex, str) -> (pd.Series, pd.DataFrame)``
|
|
make_trading_environment() -> TradingEnvironment
|
|
A class method that constructs the trading environment for the class.
|
|
If this is overridden then ``make_load_function`` or the class
|
|
attributes may not be respected.
|
|
|
|
See Also
|
|
--------
|
|
:class:`zipline.finance.trading.TradingEnvironment`
|
|
"""
|
|
TRADING_ENV_MIN_DATE = None
|
|
TRADING_ENV_MAX_DATE = None
|
|
TRADING_ENV_TRADING_CALENDAR = tradingcalendar
|
|
|
|
@classmethod
|
|
def make_load_function(cls):
|
|
return None
|
|
|
|
@classmethod
|
|
def make_trading_environment(cls):
|
|
return TradingEnvironment(
|
|
load=cls.make_load_function(),
|
|
asset_db_path=cls.asset_finder.engine,
|
|
min_date=cls.TRADING_ENV_MIN_DATE,
|
|
max_date=cls.TRADING_ENV_MAX_DATE,
|
|
env_trading_calendar=cls.TRADING_ENV_TRADING_CALENDAR,
|
|
)
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithTradingEnvironment, cls).init_class_fixtures()
|
|
cls.env = cls.make_trading_environment()
|
|
|
|
|
|
class WithSimParams(WithTradingEnvironment):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.sim_params as a class level fixture.
|
|
|
|
The arguments used to construct the trading environment may be overridded
|
|
by putting ``SIM_PARAMS_{argname}`` in the class dict except for the
|
|
trading environment which is overridden with the mechanisms provided by
|
|
``WithTradingEnvironment``.
|
|
|
|
Attributes
|
|
----------
|
|
SIM_PARAMS_YEAR : int
|
|
SIM_PARAMS_CAPITAL_BASE : float
|
|
SIM_PARAMS_NUM_DAYS : int
|
|
SIM_PARAMS_DATA_FREQUENCY : {'daily', 'minute'}
|
|
SIM_PARAMS_EMISSION_RATE : {'daily', 'minute'}
|
|
Forwarded to ``factory.create_simulation_parameters``.
|
|
|
|
SIM_PARAMS_START : datetime
|
|
SIM_PARAMS_END : datetime
|
|
Forwarded to ``factory.create_simulation_parameters``. If not
|
|
explicitly overridden these will be ``START_DATE`` and ``END_DATE``
|
|
|
|
See Also
|
|
--------
|
|
zipline.utils.factory.create_simulation_parameters
|
|
"""
|
|
SIM_PARAMS_YEAR = None
|
|
SIM_PARAMS_CAPITAL_BASE = 1.0e5
|
|
SIM_PARAMS_NUM_DAYS = None
|
|
SIM_PARAMS_DATA_FREQUENCY = 'daily'
|
|
SIM_PARAMS_EMISSION_RATE = 'daily'
|
|
|
|
SIM_PARAMS_START = alias('START_DATE')
|
|
SIM_PARAMS_END = alias('END_DATE')
|
|
|
|
@classmethod
|
|
def make_simparams(cls):
|
|
return factory.create_simulation_parameters(
|
|
year=cls.SIM_PARAMS_YEAR,
|
|
start=cls.SIM_PARAMS_START,
|
|
end=cls.SIM_PARAMS_END,
|
|
num_days=cls.SIM_PARAMS_NUM_DAYS,
|
|
capital_base=cls.SIM_PARAMS_CAPITAL_BASE,
|
|
data_frequency=cls.SIM_PARAMS_DATA_FREQUENCY,
|
|
emission_rate=cls.SIM_PARAMS_EMISSION_RATE,
|
|
env=cls.env,
|
|
)
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithSimParams, cls).init_class_fixtures()
|
|
cls.sim_params = cls.make_simparams()
|
|
|
|
|
|
class WithNYSETradingDays(object):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.trading_days as a class-level fixture.
|
|
|
|
After init_class_fixtures has been called, `cls.trading_days` is populated
|
|
with a DatetimeIndex containing NYSE calendar trading days ranging from:
|
|
|
|
(DATA_MAX_DAY - (cls.TRADING_DAY_COUNT) -> DATA_MAX_DAY)
|
|
|
|
Attributes
|
|
----------
|
|
DATA_MAX_DAY : datetime
|
|
The most recent trading day in the calendar.
|
|
TRADING_DAY_COUNT : int
|
|
The number of days to put in the calendar. The default value of
|
|
``TRADING_DAY_COUNT`` is 126 (half a trading-year). Inheritors can
|
|
override TRADING_DAY_COUNT to request more or less data.
|
|
"""
|
|
DATA_MIN_DAY = alias('START_DATE')
|
|
DATA_MAX_DAY = alias('END_DATE')
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithNYSETradingDays, cls).init_class_fixtures()
|
|
|
|
all_days = tradingcalendar.trading_days
|
|
start_loc = all_days.get_loc(cls.DATA_MIN_DAY, 'bfill')
|
|
end_loc = all_days.get_loc(cls.DATA_MAX_DAY, 'ffill')
|
|
|
|
cls.trading_days = all_days[start_loc:end_loc + 1]
|
|
|
|
|
|
class WithTmpDir(object):
|
|
"""
|
|
ZiplineTestCase mixing providing cls.tmpdir as a class-level fixture.
|
|
|
|
After init_class_fixtures has been called, `cls.tmpdir` is populated with
|
|
a `testfixtures.TempDirectory` object whose path is `cls.TMP_DIR_PATH`.
|
|
|
|
Attributes
|
|
----------
|
|
TMP_DIR_PATH : str
|
|
The path to the new directory to create. By default this is None
|
|
which will create a unique directory in /tmp.
|
|
"""
|
|
TMP_DIR_PATH = None
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithTmpDir, cls).init_class_fixtures()
|
|
cls.tmpdir = cls.enter_class_context(
|
|
tmp_dir(path=cls.TMP_DIR_PATH),
|
|
)
|
|
|
|
|
|
class WithInstanceTmpDir(object):
|
|
"""
|
|
ZiplineTestCase mixing providing self.tmpdir as an instance-level fixture.
|
|
|
|
After init_instance_fixtures has been called, `self.tmpdir` is populated
|
|
with a `testfixtures.TempDirectory` object whose path is
|
|
`cls.TMP_DIR_PATH`.
|
|
|
|
Attributes
|
|
----------
|
|
INSTANCE_TMP_DIR_PATH : str
|
|
The path to the new directory to create. By default this is None
|
|
which will create a unique directory in /tmp.
|
|
"""
|
|
INSTANCE_TMP_DIR_PATH = None
|
|
|
|
def init_instance_fixtures(self):
|
|
super(WithInstanceTmpDir, self).init_instance_fixtures()
|
|
self.instance_tmpdir = self.enter_instance_context(
|
|
tmp_dir(path=self.INSTANCE_TMP_DIR_PATH),
|
|
)
|
|
|
|
|
|
class WithBcolzDailyBarReader(WithTradingEnvironment, WithTmpDir):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.bcolz_daily_bar_path,
|
|
cls.bcolz_daily_bar_ctable, and cls.bcolz_daily_bar_reader class level
|
|
fixtures.
|
|
|
|
After init_class_fixtures has been called:
|
|
- `cls.bcolz_daily_bar_path` is populated with
|
|
`cls.tmpdir.getpath(cls.BCOLZ_DAILY_BAR_PATH)`.
|
|
- `cls.bcolz_daily_bar_ctable` is populated with data returned from
|
|
`cls.make_daily_bar_data`. By default this calls
|
|
:func:`zipline.pipeline.loaders.synthetic.make_daily_bar_data`.
|
|
- `cls.bcolz_daily_bar_reader` is a daily bar reader pointing to the
|
|
directory that was just written to.
|
|
|
|
Attributes
|
|
----------
|
|
BCOLZ_DAILY_BAR_PATH : str
|
|
The path inside the tmpdir where this will be written.
|
|
BCOLZ_DAILY_BAR_LOOKBACK_DAYS : int
|
|
The number of days of data to add before the first day. This is used
|
|
when a test needs to use history, in which case this should be set to
|
|
the largest history window that will be
|
|
requested.
|
|
BCOLZ_DAILY_BAR_USE_FULL_CALENDAR : bool
|
|
If this flag is set the ``bcolz_daily_bar_days`` will be the full
|
|
set of trading days from the trading environment. This flag overrides
|
|
``BCOLZ_DAILY_BAR_LOOKBACK_DAYS``.
|
|
BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD : int
|
|
If this flag is set, use the value as the `read_all_threshold`
|
|
parameter to BcolzDailyBarReader, otherwise use the default value.
|
|
|
|
Methods
|
|
-------
|
|
make_daily_bar_data() -> iterable[(int, pd.DataFrame)]
|
|
A class method that returns an iterator of (sid, dataframe) pairs
|
|
which will be written to the bcolz files that the class's
|
|
``BcolzDailyBarReader`` will read from. By default this creates
|
|
some simple sythetic data with
|
|
:func:`~zipline.testing.create_daily_bar_data`
|
|
|
|
See Also
|
|
--------
|
|
WithBcolzMinuteBarReader
|
|
WithDataPortal
|
|
zipline.testing.create_daily_bar_data
|
|
"""
|
|
BCOLZ_DAILY_BAR_PATH = 'daily_equity_pricing.bcolz'
|
|
BCOLZ_DAILY_BAR_LOOKBACK_DAYS = 0
|
|
BCOLZ_DAILY_BAR_USE_FULL_CALENDAR = False
|
|
BCOLZ_DAILY_BAR_START_DATE = alias('START_DATE')
|
|
BCOLZ_DAILY_BAR_END_DATE = alias('END_DATE')
|
|
BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD = None
|
|
# allows WithBcolzDailyBarReaderFromCSVs to call the `write_csvs` method
|
|
# without needing to reimplement `init_class_fixtures`
|
|
_write_method_name = 'write'
|
|
|
|
@classmethod
|
|
def make_daily_bar_data(cls):
|
|
return create_daily_bar_data(
|
|
cls.bcolz_daily_bar_days,
|
|
cls.asset_finder.sids,
|
|
)
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithBcolzDailyBarReader, cls).init_class_fixtures()
|
|
cls.bcolz_daily_bar_path = p = cls.tmpdir.makedir(
|
|
cls.BCOLZ_DAILY_BAR_PATH,
|
|
)
|
|
if cls.BCOLZ_DAILY_BAR_USE_FULL_CALENDAR:
|
|
days = cls.env.trading_days
|
|
else:
|
|
days = cls.env.days_in_range(
|
|
cls.env.trading_days[
|
|
cls.env.get_index(cls.BCOLZ_DAILY_BAR_START_DATE) -
|
|
cls.BCOLZ_DAILY_BAR_LOOKBACK_DAYS
|
|
],
|
|
cls.BCOLZ_DAILY_BAR_END_DATE,
|
|
)
|
|
cls.bcolz_daily_bar_days = days
|
|
cls.bcolz_daily_bar_ctable = t = getattr(
|
|
BcolzDailyBarWriter(p, days),
|
|
cls._write_method_name,
|
|
)(cls.make_daily_bar_data())
|
|
|
|
if cls.BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD is not None:
|
|
cls.bcolz_daily_bar_reader = BcolzDailyBarReader(
|
|
t, cls.BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD)
|
|
else:
|
|
cls.bcolz_daily_bar_reader = BcolzDailyBarReader(t)
|
|
|
|
|
|
class WithBcolzDailyBarReaderFromCSVs(WithBcolzDailyBarReader):
|
|
"""
|
|
ZiplineTestCase mixin that provides cls.bcolz_daily_bar_reader from a
|
|
mapping of sids to CSV file paths.
|
|
"""
|
|
_write_method_name = 'write_csvs'
|
|
|
|
|
|
class WithBcolzMinuteBarReader(WithTradingEnvironment, WithTmpDir):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.bcolz_minute_bar_path,
|
|
cls.bcolz_minute_bar_ctable, and cls.bcolz_minute_bar_reader class level
|
|
fixtures.
|
|
|
|
After init_class_fixtures has been called:
|
|
- `cls.bcolz_minute_bar_path` is populated with
|
|
`cls.tmpdir.getpath(cls.BCOLZ_MINUTE_BAR_PATH)`.
|
|
- `cls.bcolz_minute_bar_ctable` is populated with data returned from
|
|
`cls.make_minute_bar_data`. By default this calls
|
|
:func:`zipline.pipeline.loaders.synthetic.make_minute_bar_data`.
|
|
- `cls.bcolz_minute_bar_reader` is a minute bar reader pointing to the
|
|
directory that was just written to.
|
|
|
|
Attributes
|
|
----------
|
|
BCOLZ_MINUTE_BAR_PATH : str
|
|
The path inside the tmpdir where this will be written.
|
|
BCOLZ_MINUTE_BAR_LOOKBACK_DAYS : int
|
|
The number of days of data to add before the first day.
|
|
This is used when a test needs to use history, in which case this
|
|
should be set to the largest history window that will be requested.
|
|
BCOLZ_MINUTE_BAR_USE_FULL_CALENDAR : bool
|
|
If this flag is set the ``bcolz_daily_bar_days`` will be the full
|
|
set of trading days from the trading environment. This flag overrides
|
|
``BCOLZ_MINUTE_BAR_LOOKBACK_DAYS``.
|
|
|
|
Methods
|
|
-------
|
|
make_minute_bar_data() -> dict[int -> pd.DataFrame]
|
|
A class method that returns a dict mapping sid to dataframe
|
|
which will be written to the bcolz files that the class's
|
|
``BcolzMinuteBarReader`` will read from. By default this creates
|
|
some simple sythetic data with
|
|
:func:`~zipline.testing.create_minute_bar_data`
|
|
|
|
See Also
|
|
--------
|
|
WithBcolzDailyBarReader
|
|
WithDataPortal
|
|
zipline.testing.create_minute_bar_data
|
|
"""
|
|
BCOLZ_MINUTE_BAR_PATH = 'minute_equity_pricing.bcolz'
|
|
BCOLZ_MINUTE_BAR_LOOKBACK_DAYS = 0
|
|
BCOLZ_MINUTE_BAR_USE_FULL_CALENDAR = False
|
|
BCOLZ_MINUTE_BAR_START_DATE = alias('START_DATE')
|
|
BCOLZ_MINUTE_BAR_END_DATE = alias('END_DATE')
|
|
|
|
@classmethod
|
|
def make_minute_bar_data(cls):
|
|
return create_minute_bar_data(
|
|
cls.env.minutes_for_days_in_range(
|
|
cls.bcolz_minute_bar_days[0],
|
|
cls.bcolz_minute_bar_days[-1],
|
|
),
|
|
cls.asset_finder.sids,
|
|
)
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithBcolzMinuteBarReader, cls).init_class_fixtures()
|
|
cls.bcolz_minute_bar_path = p = cls.tmpdir.makedir(
|
|
cls.BCOLZ_MINUTE_BAR_PATH,
|
|
)
|
|
if cls.BCOLZ_MINUTE_BAR_USE_FULL_CALENDAR:
|
|
days = cls.env.trading_days
|
|
else:
|
|
days = cls.env.days_in_range(
|
|
cls.env.trading_days[
|
|
cls.env.get_index(cls.BCOLZ_MINUTE_BAR_START_DATE) -
|
|
cls.BCOLZ_MINUTE_BAR_LOOKBACK_DAYS
|
|
],
|
|
cls.BCOLZ_MINUTE_BAR_END_DATE,
|
|
)
|
|
cls.bcolz_minute_bar_days = days
|
|
writer = BcolzMinuteBarWriter(
|
|
days[0],
|
|
p,
|
|
cls.env.open_and_closes.market_open.loc[days],
|
|
cls.env.open_and_closes.market_close.loc[days],
|
|
US_EQUITIES_MINUTES_PER_DAY
|
|
)
|
|
cls.bcolz_minute_bar_data = cls.make_minute_bar_data()
|
|
for sid, df in iteritems(cls.bcolz_minute_bar_data):
|
|
writer.write(sid, df)
|
|
|
|
cls.bcolz_minute_bar_reader = BcolzMinuteBarReader(p)
|
|
|
|
|
|
class WithAdjustmentReader(WithBcolzDailyBarReader):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.adjustment_reader as a class level
|
|
fixture.
|
|
|
|
After init_class_fixtures has been called, `cls.adjustment_reader` will be
|
|
populated with a new SQLiteAdjustmentReader object. The data that will be
|
|
written can be passed by overriding `make_{field}_data` where field may
|
|
be `splits`, `mergers` `dividends`, or `stock_dividends`.
|
|
The daily bar reader used for this adjustment reader may be customized
|
|
by overriding `make_adjustment_writer_daily_bar_reader`. This is useful
|
|
to providing a `MockDailyBarReader`.
|
|
|
|
Methods
|
|
-------
|
|
make_splits_data() -> pd.DataFrame
|
|
A class method that returns a dataframe of splits data to write to the
|
|
class's adjustment db. By default this is empty.
|
|
make_mergers_data() -> pd.DataFrame
|
|
A class method that returns a dataframe of mergers data to write to the
|
|
class's adjustment db. By default this is empty.
|
|
make_dividends_data() -> pd.DataFrame
|
|
A class method that returns a dataframe of dividends data to write to
|
|
the class's adjustment db. By default this is empty.
|
|
make_stock_dividends_data() -> pd.DataFrame
|
|
A class method that returns a dataframe of stock dividends data to
|
|
write to the class's adjustment db. By default this is empty.
|
|
make_adjustment_writer_daily_bar_reader() -> pd.DataFrame
|
|
A class method that returns the daily bar reader to use for the class's
|
|
adjustment writer. By default this is the class's actual
|
|
``bcolz_daily_bar_reader`` as inherited from
|
|
``WithBcolzDailyBarReader``. This should probably not be overridden;
|
|
however, some tests used a ``MockDailyBarReader`` for this.
|
|
make_adjustment_writer(conn: sqlite3.Connection) -> AdjustmentWriter
|
|
A class method that constructs the adjustment which will be used
|
|
to write the data into the connection to be used by the class's
|
|
adjustment reader.
|
|
|
|
See Also
|
|
--------
|
|
zipline.testing.MockDailyBarReader
|
|
"""
|
|
@classmethod
|
|
def _make_data(cls):
|
|
return None
|
|
|
|
make_splits_data = _make_data
|
|
make_mergers_data = _make_data
|
|
make_dividends_data = _make_data
|
|
make_stock_dividends_data = _make_data
|
|
|
|
del _make_data
|
|
|
|
@classmethod
|
|
def make_adjustment_writer(cls, conn):
|
|
return SQLiteAdjustmentWriter(
|
|
conn,
|
|
cls.make_adjustment_writer_daily_bar_reader(),
|
|
cls.bcolz_daily_bar_days,
|
|
)
|
|
|
|
@classmethod
|
|
def make_adjustment_writer_daily_bar_reader(cls):
|
|
return cls.bcolz_daily_bar_reader
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithAdjustmentReader, cls).init_class_fixtures()
|
|
conn = sqlite3.connect(':memory:')
|
|
cls.make_adjustment_writer(conn).write(
|
|
splits=cls.make_splits_data(),
|
|
mergers=cls.make_mergers_data(),
|
|
dividends=cls.make_dividends_data(),
|
|
stock_dividends=cls.make_stock_dividends_data(),
|
|
)
|
|
cls.adjustment_reader = SQLiteAdjustmentReader(conn)
|
|
|
|
|
|
class WithPipelineEventDataLoader(with_metaclass(
|
|
compose_types(ABCMeta, type(ZiplineTestCase)), WithAssetFinder)):
|
|
"""
|
|
ZiplineTestCase mixin providing common test methods/behaviors for event
|
|
data loaders.
|
|
|
|
Attributes
|
|
----------
|
|
loader_type : PipelineLoader
|
|
The type of loader to use. This must be overridden by subclasses.
|
|
|
|
Methods
|
|
-------
|
|
get_sids() -> iterable[int]
|
|
Class method which returns the sids that need to be available to the
|
|
tests.
|
|
get_dataset() -> dict[int -> pd.DataFrmae]
|
|
Class method which returns a mapping from sid to data for that sid.
|
|
By default this is empty for every sid.
|
|
pipeline_event_loader_args(dates: pd.DatetimeIndex) -> tuple[any]
|
|
The arguments to pass to the ``loader_type`` to construct the pipeline
|
|
loader for this test.
|
|
"""
|
|
@classmethod
|
|
def get_sids(cls):
|
|
return range(0, 5)
|
|
|
|
@classmethod
|
|
def get_dataset(cls):
|
|
return {sid: pd.DataFrame() for sid in cls.get_sids()}
|
|
|
|
@abstractproperty
|
|
def loader_type(self):
|
|
raise NotImplementedError('loader_type')
|
|
|
|
@classmethod
|
|
def make_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'),
|
|
)
|
|
|
|
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)
|
|
|
|
def get_sids_to_frames(self,
|
|
zip_date_index_with_vals,
|
|
vals,
|
|
date_intervals,
|
|
dates):
|
|
"""
|
|
Construct a DataFrame that maps sid to the expected values for the
|
|
given dates.
|
|
|
|
Parameters
|
|
----------
|
|
zip_date_index_with_vals: callable
|
|
A function that returns a series of `vals` repeated based on the
|
|
number of days in the date interval for each val, indexed by the
|
|
dates in `dates`.
|
|
vals: iterable
|
|
An iterable with values that correspond to each interval in
|
|
`date_intervals`.
|
|
date_intervals: list
|
|
A list of date intervals for each sid that correspond to values in
|
|
`vals`.
|
|
dates: DatetimeIndex
|
|
The dates which will serve as the index for each Series for each
|
|
sid in the DataFrame.
|
|
"""
|
|
frame = pd.DataFrame({sid: get_values_for_date_ranges(
|
|
zip_date_index_with_vals,
|
|
vals[sid],
|
|
pd.DatetimeIndex(list(zip(*date_intervals[sid]))[0]),
|
|
pd.DatetimeIndex(list(zip(*date_intervals[sid]))[1]),
|
|
dates
|
|
) for sid in self.get_sids()[:-1]})
|
|
frame[self.get_sids()[-1]] = zip_date_index_with_vals(
|
|
dates, ['NaN'] * len(dates)
|
|
)
|
|
return frame
|
|
|
|
@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].unstack(1)[sid],
|
|
cols[col_name][sid],
|
|
check_names=False)
|
|
|
|
|
|
class WithDataPortal(WithBcolzMinuteBarReader, WithAdjustmentReader):
|
|
"""
|
|
ZiplineTestCase mixin providing self.data_portal as an instance level
|
|
fixture.
|
|
|
|
After init_instance_fixtures has been called, `self.data_portal` will be
|
|
populated with a new data portal created by passing in the class's
|
|
trading env, `cls.bcolz_minute_bar_reader`, `cls.bcolz_daily_bar_reader`,
|
|
and `cls.adjustment_reader`.
|
|
|
|
Attributes
|
|
----------
|
|
DATA_PORTAL_USE_DAILY_DATA : bool
|
|
Should the daily bar reader be used? Defaults to True.
|
|
DATA_PORTAL_USE_MINUTE_DATA : bool
|
|
Should the minute bar reader be used? Defaults to True.
|
|
DATA_PORTAL_USE_ADJUSTMENTS : bool
|
|
Should the adjustment reader be used? Defaults to True.
|
|
|
|
Methods
|
|
-------
|
|
make_data_portal() -> DataPortal
|
|
Method which returns the data portal to be used for each test case.
|
|
If this is overridden, the ``DATA_PORTAL_USE_*`` attributes may not
|
|
be respected.
|
|
"""
|
|
DATA_PORTAL_USE_DAILY_DATA = True
|
|
DATA_PORTAL_USE_MINUTE_DATA = True
|
|
DATA_PORTAL_USE_ADJUSTMENTS = True
|
|
|
|
def make_data_portal(self):
|
|
return DataPortal(
|
|
self.env,
|
|
equity_daily_reader=(
|
|
self.bcolz_daily_bar_reader
|
|
if self.DATA_PORTAL_USE_DAILY_DATA else
|
|
None
|
|
),
|
|
equity_minute_reader=(
|
|
self.bcolz_minute_bar_reader
|
|
if self.DATA_PORTAL_USE_MINUTE_DATA else
|
|
None
|
|
),
|
|
adjustment_reader=(
|
|
self.adjustment_reader
|
|
if self.DATA_PORTAL_USE_ADJUSTMENTS else
|
|
None
|
|
),
|
|
)
|
|
|
|
def init_instance_fixtures(self):
|
|
super(WithDataPortal, self).init_instance_fixtures()
|
|
self.data_portal = self.make_data_portal()
|
|
|
|
|
|
class WithResponses(object):
|
|
"""
|
|
ZiplineTestCase mixin that provides self.responses as an instance
|
|
fixture.
|
|
|
|
After init_instance_fixtures has been called, `self.responses` will be
|
|
a new `responses.RequestsMock` object. Users may add new endpoints to this
|
|
with the `self.responses.add` method.
|
|
"""
|
|
def init_instance_fixtures(self):
|
|
super(WithResponses, self).init_instance_fixtures()
|
|
self.responses = self.enter_instance_context(
|
|
responses.RequestsMock(),
|
|
)
|
|
|
|
|
|
class WithNextAndPreviousEventDataLoader(WithPipelineEventDataLoader):
|
|
"""
|
|
ZiplineTestCase mixin extending common functionality for event data
|
|
loader tests that have both next and previous events.
|
|
|
|
`base_cases` should be used as the template to test cases that combine
|
|
knowledge date (timestamp) and some 'other_date' in various ways.
|
|
`next_date_intervals` gives the date intervals for the next event based
|
|
on the dates given in `base_cases`.
|
|
`next_dates` gives the next date from `other_date` which is known about at
|
|
each interval.
|
|
`prev_date_intervals` gives the date intervals for each sid for the
|
|
previous event based on the dates given in `base_cases`.
|
|
`prev_dates` gives the previous date from `other_date` which is known
|
|
about at each interval.
|
|
`get_expected_previous_event_dates` is a convenience function that fills
|
|
a DataFrame with the previously known dates for each sid for the given
|
|
dates.
|
|
`get_expected_next_event_dates` is a convenience function that fills
|
|
a DataFrame with the next known dates for each sid for the given
|
|
dates.
|
|
"""
|
|
base_cases = [
|
|
# K1--K2--A1--A2.
|
|
pd.DataFrame({
|
|
TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']),
|
|
'other_date': 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']),
|
|
'other_date': 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']),
|
|
'other_date': pd.to_datetime(['2014-01-10', '2014-01-20']),
|
|
}),
|
|
# K1 == K2.
|
|
pd.DataFrame({
|
|
TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2),
|
|
'other_date': pd.to_datetime(['2014-01-10', '2014-01-15']),
|
|
}),
|
|
pd.DataFrame(
|
|
columns=['other_date',
|
|
TS_FIELD_NAME],
|
|
dtype='datetime64[ns]'
|
|
),
|
|
]
|
|
|
|
next_date_intervals = [
|
|
[['2014-01-01', '2014-01-04'],
|
|
['2014-01-05', '2014-01-15'],
|
|
['2014-01-16', '2014-01-20'],
|
|
['2014-01-21', '2014-01-31']],
|
|
[['2014-01-01', '2014-01-04'],
|
|
['2014-01-05', '2014-01-09'],
|
|
['2014-01-10', '2014-01-15'],
|
|
['2014-01-16', '2014-01-20'],
|
|
['2014-01-21', '2014-01-31']],
|
|
[['2014-01-01', '2014-01-04'],
|
|
['2014-01-05', '2014-01-10'],
|
|
['2014-01-11', '2014-01-14'],
|
|
['2014-01-15', '2014-01-20'],
|
|
['2014-01-21', '2014-01-31']],
|
|
[['2014-01-01', '2014-01-04'],
|
|
['2014-01-05', '2014-01-10'],
|
|
['2014-01-11', '2014-01-15'],
|
|
['2014-01-16', '2014-01-31']]
|
|
]
|
|
|
|
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 = [
|
|
[['2014-01-01', '2014-01-14'],
|
|
['2014-01-15', '2014-01-19'],
|
|
['2014-01-20', '2014-01-31']],
|
|
[['2014-01-01', '2014-01-14'],
|
|
['2014-01-15', '2014-01-19'],
|
|
['2014-01-20', '2014-01-31']],
|
|
[['2014-01-01', '2014-01-09'],
|
|
['2014-01-10', '2014-01-19'],
|
|
['2014-01-20', '2014-01-31']],
|
|
[['2014-01-01', '2014-01-09'],
|
|
['2014-01-10', '2014-01-14'],
|
|
['2014-01-15', '2014-01-31']]
|
|
]
|
|
|
|
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']
|
|
]
|
|
|
|
def get_expected_previous_event_dates(self, dates):
|
|
return self.get_sids_to_frames(
|
|
zip_with_dates,
|
|
self.prev_dates,
|
|
self.prev_date_intervals,
|
|
dates
|
|
)
|
|
|
|
def get_expected_next_event_dates(self, dates):
|
|
return self.get_sids_to_frames(
|
|
zip_with_dates,
|
|
self.next_dates,
|
|
self.next_date_intervals,
|
|
dates
|
|
)
|