Files
catalyst/zipline/testing/fixtures.py
T
Eddie Hebert f4891b0a08 TST: Key trading calendar fixture with Asset types
Instead of using strings of 'equities' and 'futures', use the Asset
subclasses to key the trading calendar fixtures.
2016-08-08 03:49:48 -04:00

1186 lines
42 KiB
Python

import sqlite3
from unittest import TestCase
from contextlib2 import ExitStack
from logbook import NullHandler, Logger
from six import with_metaclass, iteritems
from toolz import flip
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.resample import minute_to_session
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 factory
from ..utils.classproperty import classproperty
from ..utils.final import FinalMeta, final
from .core import tmp_asset_finder, make_simple_equity_info
from zipline.assets import Equity
from zipline.pipeline import SimplePipelineEngine
from zipline.pipeline.loaders.testing import make_seeded_random_loader
from zipline.utils.calendars import (
get_calendar,
register_calendar)
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):
# We need to get this before it's deleted by the loop.
stack = cls._class_teardown_stack
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)
stack.close()
@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):
# We need to get this before it's deleted by the loop.
stack = self._instance_teardown_stack
for attr in set(vars(self)) - self._pre_setup_attrs:
delattr(self, attr)
stack.close()
@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(C):
... 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_db_url() -> string
A class method which returns the URL at which to create the SQLAlchemy
engine. By default provides a URL for an in-memory database.
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):
register_calendar("TEST", get_calendar("NYSE"), force=True)
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_db_url(cls):
return 'sqlite:///:memory:'
@classmethod
def make_asset_finder(cls):
return cls.enter_class_context(tmp_asset_finder(
url=cls.make_asset_finder_db_url(),
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 WithTradingCalendars(object):
"""
ZiplineTestCase mixin providing cls.trading_calendar,
cls.all_trading_calendars, cls.trading_calendar_for_asset_type as a
class-level fixture.
After ``init_class_fixtures`` has been called:
- `cls.trading_calendar` is populated with a default of the nyse trading
calendar for compatibility with existing tests
- `cls.all_trading_calendars` is populated with the trading calendars
keyed by name,
- `cls.trading_calendar_for_asset_type` is populated with the trading
calendars keyed by the asset type which uses the respective calendar.
Attributes
----------
TRADING_CALENDAR_STRS : iterable
iterable of identifiers of the calendars to use.
TRADING_CALENDAR_FOR_ASSET_TYPE : dict
A dictionay which maps asset type names to the calendar associated
with that asset type.
"""
TRADING_CALENDAR_STRS = ('NYSE',)
TRADING_CALENDAR_FOR_ASSET_TYPE = {Equity: 'NYSE'}
# For backwards compatibility, exisitng tests and fixtures refer to
# `trading_calendar` with the assumption that the value is the NYSE
# calendar.
trading_calendar = alias('nyse_calendar')
@classmethod
def init_class_fixtures(cls):
super(WithTradingCalendars, cls).init_class_fixtures()
cls.trading_calendars = {}
for cal_str in cls.TRADING_CALENDAR_STRS:
# Set name to allow aliasing.
calendar = get_calendar(cal_str)
setattr(cls,
'{0}_calendar'.format(cal_str.lower()), calendar)
cls.trading_calendars[cal_str] = calendar
for asset_type, cal_str in iteritems(
cls.TRADING_CALENDAR_FOR_ASSET_TYPE):
calendar = get_calendar(cal_str)
cls.trading_calendars[asset_type] = calendar
class WithTradingEnvironment(WithAssetFinder, WithTradingCalendars):
"""
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`
"""
@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,
trading_calendar=cls.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,
trading_calendar=cls.trading_calendar,
)
@classmethod
def init_class_fixtures(cls):
super(WithSimParams, cls).init_class_fixtures()
cls.sim_params = cls.make_simparams()
class WithTradingSessions(WithTradingCalendars):
"""
ZiplineTestCase mixin providing cls.trading_days, cls.all_trading_sessions
as a class-level fixture.
After init_class_fixtures has been called, `cls.all_trading_sessions`
is populated with a dictionary of calendar name to the DatetimeIndex
containing the calendar trading days ranging from:
(DATA_MAX_DAY - (cls.TRADING_DAY_COUNT) -> DATA_MAX_DAY)
`cls.trading_days`, for compatibility with existing tests which make the
assumption that trading days are equity only, defaults to the nyse trading
sessions.
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')
# For backwards compatibility, exisitng tests and fixtures refer to
# `trading_days` with the assumption that the value is days of the NYSE
# calendar.
trading_days = alias('nyse_sessions')
@classmethod
def init_class_fixtures(cls):
super(WithTradingSessions, cls).init_class_fixtures()
cls.trading_sessions = {}
for cal_str in cls.TRADING_CALENDAR_STRS:
trading_calendar = cls.trading_calendars[cal_str]
all_sessions = trading_calendar.all_sessions
start_loc = all_sessions.get_loc(cls.DATA_MIN_DAY, 'bfill')
end_loc = all_sessions.get_loc(cls.DATA_MAX_DAY, 'ffill')
sessions = all_sessions[start_loc:end_loc + 1]
# Set name for aliasing.
setattr(cls,
'{0}_sessions'.format(cal_str.lower()), sessions)
cls.trading_sessions[cal_str] = sessions
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 WithEquityDailyBarData(WithTradingEnvironment):
"""
ZiplineTestCase mixin providing cls.make_equity_daily_bar_data.
Attributes
----------
EQUITY_DAILY_BAR_START_DATE : Timestamp
The date at to which to start creating data. This defaults to
``START_DATE``.
EQUITY_DAILY_BAR_END_DATE = Timestamp
The end date up to which to create data. This defaults to ``END_DATE``.
EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE : bool
If this flag is set, `make_equity_daily_bar_data` will read data from
the minute bars defined by `WithMinuteBarData`.
The current default is `False`, but could be `True` in the future.
Methods
-------
make_equity_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
--------
WithEquityMinuteBarData
zipline.testing.create_daily_bar_data
"""
EQUITY_DAILY_BAR_LOOKBACK_DAYS = 0
EQUITY_DAILY_BAR_USE_FULL_CALENDAR = False
EQUITY_DAILY_BAR_START_DATE = alias('START_DATE')
EQUITY_DAILY_BAR_END_DATE = alias('END_DATE')
EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE = None
@classmethod
def _make_equity_daily_bar_from_minute(cls):
assets = cls.asset_finder.retrieve_all(cls.asset_finder.sids)
minute_data = dict(cls.make_equity_minute_bar_data())
for asset in assets:
yield asset.sid, minute_to_session(minute_data[asset.sid],
cls.trading_calendar)
@classmethod
def make_equity_daily_bar_data(cls):
# Requires a WithEquityMinuteBarData to come before in the MRO.
# Resample that data so that daily and minute bar data are aligned.
if cls.EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE:
return cls._make_equity_daily_bar_from_minute()
else:
return create_daily_bar_data(
cls.equity_daily_bar_days,
cls.asset_finder.sids,
)
@classmethod
def init_class_fixtures(cls):
super(WithEquityDailyBarData, cls).init_class_fixtures()
if cls.EQUITY_DAILY_BAR_USE_FULL_CALENDAR:
days = cls.trading_calendar.all_sessions
else:
if cls.trading_calendar.is_session(
cls.EQUITY_DAILY_BAR_START_DATE
):
first_session = cls.EQUITY_DAILY_BAR_START_DATE
else:
first_session = cls.trading_calendar.minute_to_session_label(
pd.Timestamp(cls.EQUITY_DAILY_BAR_START_DATE)
)
if cls.EQUITY_DAILY_BAR_LOOKBACK_DAYS > 0:
first_session = cls.trading_calendar.sessions_window(
first_session,
-1 * cls.EQUITY_DAILY_BAR_LOOKBACK_DAYS
)[0]
days = cls.trading_calendar.sessions_in_range(
first_session,
cls.EQUITY_DAILY_BAR_END_DATE,
)
cls.equity_daily_bar_days = days
class WithBcolzEquityDailyBarReader(WithEquityDailyBarData, WithTmpDir):
"""
ZiplineTestCase mixin providing cls.bcolz_daily_bar_path,
cls.bcolz_daily_bar_ctable, and cls.bcolz_equity_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_equity_daily_bar_data`. By default this calls
:func:`zipline.pipeline.loaders.synthetic.make_equity_daily_bar_data`.
- `cls.bcolz_equity_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.
EQUITY_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.
EQUITY_DAILY_BAR_USE_FULL_CALENDAR : bool
If this flag is set the ``equity_daily_bar_days`` will be the full
set of trading days from the trading environment. This flag overrides
``EQUITY_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.
EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE : bool
If this flag is set, `make_equity_daily_bar_data` will read data from
the minute bar reader defined by a `WithBcolzEquityMinuteBarReader`.
Methods
-------
make_bcolz_daily_bar_rootdir_path() -> string
A class method that returns the path for the rootdir of the daily
bars ctable. By default this is a subdirectory BCOLZ_DAILY_BAR_PATH in
the shared temp directory.
See Also
--------
WithBcolzEquityMinuteBarReader
WithDataPortal
zipline.testing.create_daily_bar_data
"""
BCOLZ_DAILY_BAR_PATH = 'daily_equity_pricing.bcolz'
BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD = None
EQUITY_DAILY_BAR_SOURCE_FROM_MINUTE = False
# allows WithBcolzEquityDailyBarReaderFromCSVs to call the
# `write_csvs`method without needing to reimplement `init_class_fixtures`
_write_method_name = 'write'
@classmethod
def make_bcolz_daily_bar_rootdir_path(cls):
return cls.tmpdir.makedir(cls.BCOLZ_DAILY_BAR_PATH)
@classmethod
def init_class_fixtures(cls):
super(WithBcolzEquityDailyBarReader, cls).init_class_fixtures()
cls.bcolz_daily_bar_path = p = cls.make_bcolz_daily_bar_rootdir_path()
days = cls.equity_daily_bar_days
cls.bcolz_daily_bar_ctable = t = getattr(
BcolzDailyBarWriter(p, cls.trading_calendar, days[0], days[-1]),
cls._write_method_name,
)(cls.make_equity_daily_bar_data())
if cls.BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD is not None:
cls.bcolz_equity_daily_bar_reader = BcolzDailyBarReader(
t, cls.BCOLZ_DAILY_BAR_READ_ALL_THRESHOLD)
else:
cls.bcolz_equity_daily_bar_reader = BcolzDailyBarReader(t)
class WithBcolzEquityDailyBarReaderFromCSVs(WithBcolzEquityDailyBarReader):
"""
ZiplineTestCase mixin that provides
cls.bcolz_equity_daily_bar_reader from a mapping of sids to CSV
file paths.
"""
_write_method_name = 'write_csvs'
class WithEquityMinuteBarData(WithTradingEnvironment):
"""
ZiplineTestCase mixin providing cls.equity_minute_bar_days.
After init_class_fixtures has been called:
- `cls.equyt_minute_bar_days` has the range over which data has been
generated.
Attributes
----------
EQUITY_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.
EQUITY_MINUTE_BAR_USE_FULL_CALENDAR : bool
If this flag is set the ``equity_daily_bar_days`` will be the full
set of trading days from the trading environment. This flag overrides
``EQUITY_MINUTE_BAR_LOOKBACK_DAYS``.
EQUITY_MINUTE_BAR_START_DATE : Timestamp
The date at to which to start creating data. This defaults to
``START_DATE``.
EQUITY_MINUTE_BAR_END_DATE = Timestamp
The end date up to which to create data. This defaults to ``END_DATE``.
Methods
-------
make_equity_minute_bar_data() -> iterable[(int, pd.DataFrame)]
A class method that returns a dict mapping sid to dataframe
which will be written to into the the format of the inherited
class which writes the minute bar data for use by a reader.
By default this creates some simple sythetic data with
:func:`~zipline.testing.create_minute_bar_data`
See Also
--------
WithEquityDailyBarData
zipline.testing.create_minute_bar_data
"""
EQUITY_MINUTE_BAR_LOOKBACK_DAYS = 0
EQUITY_MINUTE_BAR_USE_FULL_CALENDAR = False
EQUITY_MINUTE_BAR_START_DATE = alias('START_DATE')
EQUITY_MINUTE_BAR_END_DATE = alias('END_DATE')
@classmethod
def make_equity_minute_bar_data(cls):
trading_calendar = cls.trading_calendars[Equity]
return create_minute_bar_data(
trading_calendar.minutes_for_sessions_in_range(
cls.equity_minute_bar_days[0],
cls.equity_minute_bar_days[-1],
),
cls.asset_finder.sids,
)
@classmethod
def init_class_fixtures(cls):
super(WithEquityMinuteBarData, cls).init_class_fixtures()
trading_calendar = cls.trading_calendars[Equity]
if cls.EQUITY_MINUTE_BAR_USE_FULL_CALENDAR:
days = trading_calendar.all_execution_days
else:
first_session = trading_calendar.minute_to_session_label(
pd.Timestamp(cls.EQUITY_MINUTE_BAR_START_DATE)
)
if cls.EQUITY_MINUTE_BAR_LOOKBACK_DAYS > 0:
first_session = trading_calendar.sessions_window(
first_session,
-1 * cls.EQUITY_MINUTE_BAR_LOOKBACK_DAYS
)[0]
days = trading_calendar.sessions_in_range(
first_session,
cls.EQUITY_MINUTE_BAR_END_DATE
)
cls.equity_minute_bar_days = days
class WithBcolzEquityMinuteBarReader(WithEquityMinuteBarData, WithTmpDir):
"""
ZiplineTestCase mixin providing cls.bcolz_minute_bar_path,
cls.bcolz_minute_bar_ctable, and cls.bcolz_equity_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_equity_minute_bar_data`. By default this calls
:func:`zipline.pipeline.loaders.synthetic.make_equity_minute_bar_data`.
- `cls.bcolz_equity_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.
EQUITY_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 ``equity_daily_bar_days`` will be the full
set of trading days from the trading environment. This flag overrides
``EQUITY_MINUTE_BAR_LOOKBACK_DAYS``.
Methods
-------
make_bcolz_minute_bar_rootdir_path() -> string
A class method that returns the path for the directory that contains
the minute bar ctables. By default this is a subdirectory
BCOLZ_MINUTE_BAR_PATH in the shared temp directory.
See Also
--------
WithBcolzEquityDailyBarReader
WithDataPortal
zipline.testing.create_minute_bar_data
"""
BCOLZ_MINUTE_BAR_PATH = 'minute_equity_pricing.bcolz'
@classmethod
def make_bcolz_minute_bar_rootdir_path(cls):
return cls.tmpdir.makedir(cls.BCOLZ_MINUTE_BAR_PATH)
@classmethod
def init_class_fixtures(cls):
super(WithBcolzEquityMinuteBarReader, cls).init_class_fixtures()
cls.bcolz_minute_bar_path = p = \
cls.make_bcolz_minute_bar_rootdir_path()
days = cls.equity_minute_bar_days
writer = BcolzMinuteBarWriter(
days[0],
p,
cls.trading_calendar.schedule.market_open.loc[days],
cls.trading_calendar.schedule.market_close.loc[days],
US_EQUITIES_MINUTES_PER_DAY
)
writer.write(cls.make_equity_minute_bar_data())
cls.bcolz_equity_minute_bar_reader = \
BcolzMinuteBarReader(p)
class WithAdjustmentReader(WithBcolzEquityDailyBarReader):
"""
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_equity_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_db_conn_str() -> string
A class method that returns the sqlite3 connection string for the
database in to which the adjustments will be written. By default this
is an in-memory database.
make_adjustment_writer_equity_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_equity_daily_bar_reader`` as inherited from
``WithBcolzEquityDailyBarReader``. 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_equity_daily_bar_reader(),
cls.equity_daily_bar_days,
)
@classmethod
def make_adjustment_writer_equity_daily_bar_reader(cls):
return cls.bcolz_equity_daily_bar_reader
@classmethod
def make_adjustment_db_conn_str(cls):
return ':memory:'
@classmethod
def init_class_fixtures(cls):
super(WithAdjustmentReader, cls).init_class_fixtures()
conn = sqlite3.connect(cls.make_adjustment_db_conn_str())
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 WithSeededRandomPipelineEngine(WithTradingSessions, WithAssetFinder):
"""
ZiplineTestCase mixin providing class-level fixtures for running pipelines
against deterministically-generated random data.
Attributes
----------
SEEDED_RANDOM_PIPELINE_SEED : int
Fixture input. Random seed used to initialize the random state loader.
seeded_random_loader : SeededRandomLoader
Fixture output. Loader capable of providing columns for
zipline.pipeline.data.testing.TestingDataSet.
seeded_random_engine : SimplePipelineEngine
Fixture output. A pipeline engine that will use seeded_random_loader
as its only data provider.
Methods
-------
run_pipeline(start_date, end_date)
Run a pipeline with self.seeded_random_engine.
See Also
--------
zipline.pipeline.loaders.synthetic.SeededRandomLoader
zipline.pipeline.loaders.testing.make_seeded_random_loader
zipline.pipeline.engine.SimplePipelineEngine
"""
SEEDED_RANDOM_PIPELINE_SEED = 42
@classmethod
def init_class_fixtures(cls):
super(WithSeededRandomPipelineEngine, cls).init_class_fixtures()
cls._sids = cls.asset_finder.sids
cls.seeded_random_loader = loader = make_seeded_random_loader(
cls.SEEDED_RANDOM_PIPELINE_SEED,
cls.trading_days,
cls._sids,
)
cls.seeded_random_engine = SimplePipelineEngine(
get_loader=lambda column: loader,
calendar=cls.trading_days,
asset_finder=cls.asset_finder,
)
def raw_expected_values(self, column, start_date, end_date):
"""
Get an array containing the raw values we expect to be produced for the
given dates between start_date and end_date, inclusive.
"""
all_values = self.seeded_random_loader.values(
column.dtype,
self.trading_days,
self._sids,
)
row_slice = self.trading_days.slice_indexer(start_date, end_date)
return all_values[row_slice]
def run_pipeline(self, pipeline, start_date, end_date):
"""
Run a pipeline with self.seeded_random_engine.
"""
if start_date not in self.trading_days:
raise AssertionError("Start date not in calendar: %s" % start_date)
if end_date not in self.trading_days:
raise AssertionError("Start date not in calendar: %s" % start_date)
return self.seeded_random_engine.run_pipeline(
pipeline,
start_date,
end_date,
)
class WithDataPortal(WithAdjustmentReader,
# Ordered so that bcolz minute reader is used first.
WithBcolzEquityMinuteBarReader):
"""
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_equity_minute_bar_reader`,
`cls.bcolz_equity_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
DATA_PORTAL_FIRST_TRADING_DAY = None
def make_data_portal(self):
if self.DATA_PORTAL_FIRST_TRADING_DAY is None:
if self.DATA_PORTAL_USE_MINUTE_DATA:
self.DATA_PORTAL_FIRST_TRADING_DAY = (
self.bcolz_equity_minute_bar_reader.
first_trading_day)
elif self.DATA_PORTAL_USE_DAILY_DATA:
self.DATA_PORTAL_FIRST_TRADING_DAY = (
self.bcolz_equity_daily_bar_reader.
first_trading_day)
return DataPortal(
self.env.asset_finder,
self.trading_calendar,
first_trading_day=self.DATA_PORTAL_FIRST_TRADING_DAY,
equity_daily_reader=(
self.bcolz_equity_daily_bar_reader
if self.DATA_PORTAL_USE_DAILY_DATA else
None
),
equity_minute_reader=(
self.bcolz_equity_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(),
)