MAINT: make the data loading apis more consistent.

Changes BcolzDailyBarWriter to not be an abc, data is passed as an
iterator of (sid, dataframe) pairs to the write method.

Changes the AssetsDBWriter to be a single class which accepts an engine
at construction time and has a `write` method for writing dataframes for
the various tables. We no longer support writing the various other data
types, callers should coerce their data into a dataframe themselves. See
zipline.assets.synthetic for some helpers to do this.

Adds many new fixtures and updates some existing fixtures to use the new
ones:

WithDefaultDateBounds
  A fixture that provides the suite a START_DATE and END_DATE. This is
  meant to make it easy for other fixtures to synchronize their date
  ranges without depending on eachother in strange ways. For example,
  WithBcolzMinuteBarReader and WithBcolzDailyBarReader by default should
  both have data for the same dates, so they may use depend on
  WithDefaultDates without forcing a dependency between them.

WithTmpDir, WithInstanceTmpDir
  Provides the suite or individual test case a temporary directory.

WithBcolzDailyBarReader
  Provides the suite a BcolzDailyBarReader which reads from bcolz data
  written to a temporary directory. The data will be read from
  dataframes and then converted to bcolz files with
  BcolzDailyBarWriter.write

WithBcolzDailyBarReaderFromCSVs
  Provides the suite a BcolzDailyBarReader which reads from bcolz data
  written to a temporary directory. The data will be read from a
  collection of CSV files and then converted into the bcolz data through
  BcolzDailyBarWriter.write_csvs

WithBcolzMinuteBarReader
  Provides the suite a BcolzMinuteBarReader which reads from bcolz data
  written to a temporary directory. The data will be read from
  dataframes and then converted to bcolz files with
  BcolzMinuteBarWriter.write

WithAdjustmentReader
  Provides the suite a SQLiteAdjustmentReader which reads from an in
  memory sqlite database. The data will be read from dataframes and then
  converted into sqlite with SQLiteAdjustmentWriter.write

WithDataPortal
  Provides each test case a DataPortal object with data from temporary
  resources.
This commit is contained in:
Joe Jevnik
2016-04-15 23:46:10 -04:00
parent 8c64cc80ec
commit bc0b117dc9
68 changed files with 5145 additions and 4922 deletions
+2
View File
@@ -66,3 +66,5 @@ zipline.iml
# PyCharm custom settings
.idea
TAGS
+64
View File
@@ -0,0 +1,64 @@
Development
-----------
:Release: 1.0.0
:Date: TBD
.. warning::
This release is still under active development. All changes listed are
subject to change at any time.
Highlights
~~~~~~~~~~
None
Enhancements
~~~~~~~~~~~~
* Made the data loading classes have more consistent interfaces. This includes
the equity bar writers, adjustment writer, and asset db writer. The new
interface is that the resource to be written to is passed at construction time
and the data to write is provided later to the `write` method as a
dataframe. This model allows us to pass these writer objects around as a
resource for other classes and functions to consume (:issue:`1109`).
Experimental Features
~~~~~~~~~~~~~~~~~~~~~
.. warning::
Experimental features are subject to change.
None
Bug Fixes
~~~~~~~~~
None
Performance
~~~~~~~~~~~
None
Maintenance and Refactorings
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
None
Build
~~~~~
None
Documentation
~~~~~~~~~~~~~
None
Miscellaneous
~~~~~~~~~~~~~
None
+1
View File
@@ -51,6 +51,7 @@ click==4.0.0
# FUNctional programming utilities
toolz==0.7.4
multipledispatch==0.4.8
# Asset writer and finder
sqlalchemy==1.0.8
+64 -55
View File
@@ -12,8 +12,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import TestCase
from nose_parameterized import parameterized
from numpy import (
arange,
@@ -28,18 +26,25 @@ from pandas import (
Timestamp,
)
from pandas.util.testing import assert_index_equal
from testfixtures import TempDirectory
from zipline.pipeline.loaders.synthetic import (
SyntheticDailyBarWriter,
)
from zipline.data.us_equity_pricing import (
BcolzDailyBarReader,
NoDataOnDate
NoDataOnDate,
)
from zipline.finance.trading import TradingEnvironment
from zipline.pipeline.data import USEquityPricing
from zipline.pipeline.loaders.synthetic import (
OHLCV,
asset_start,
asset_end,
expected_daily_bar_value,
expected_daily_bar_values_2d,
make_daily_bar_data,
)
from zipline.testing import seconds_to_timestamp
from zipline.testing.fixtures import (
WithBcolzDailyBarReader,
ZiplineTestCase,
)
TEST_CALENDAR_START = Timestamp('2015-06-01', tz='UTC')
TEST_CALENDAR_STOP = Timestamp('2015-06-30', tz='UTC')
@@ -72,58 +77,57 @@ EQUITY_INFO = DataFrame(
TEST_QUERY_ASSETS = EQUITY_INFO.index
class BcolzDailyBarTestCase(TestCase):
class BcolzDailyBarTestCase(WithBcolzDailyBarReader, ZiplineTestCase):
BCOLZ_DAILY_BAR_START_DATE = TEST_CALENDAR_START
BCOLZ_DAILY_BAR_END_DATE = TEST_CALENDAR_STOP
@classmethod
def setUpClass(cls):
all_trading_days = TradingEnvironment().trading_days
def make_equity_info(cls):
return EQUITY_INFO
@classmethod
def make_daily_bar_data(cls):
return make_daily_bar_data(
EQUITY_INFO,
cls.bcolz_daily_bar_days,
)
@classmethod
def init_class_fixtures(cls):
super(BcolzDailyBarTestCase, cls).init_class_fixtures()
all_trading_days = cls.env.trading_days
cls.trading_days = all_trading_days[
all_trading_days.get_loc(TEST_CALENDAR_START):
all_trading_days.get_loc(TEST_CALENDAR_STOP) + 1
]
def setUp(self):
self.asset_info = EQUITY_INFO
self.writer = SyntheticDailyBarWriter(
self.asset_info,
self.trading_days,
)
self.dir_ = TempDirectory()
self.dir_.create()
self.dest = self.dir_.getpath('daily_equity_pricing.bcolz')
def tearDown(self):
self.dir_.cleanup()
@property
def assets(self):
return self.asset_info.index
return EQUITY_INFO.index
def trading_days_between(self, start, end):
return self.trading_days[self.trading_days.slice_indexer(start, end)]
def asset_start(self, asset_id):
return self.writer.asset_start(asset_id)
return asset_start(EQUITY_INFO, asset_id)
def asset_end(self, asset_id):
return self.writer.asset_end(asset_id)
return asset_end(EQUITY_INFO, asset_id)
def dates_for_asset(self, asset_id):
start, end = self.asset_start(asset_id), self.asset_end(asset_id)
return self.trading_days_between(start, end)
def test_write_ohlcv_content(self):
result = self.writer.write(self.dest, self.trading_days, self.assets)
for column in SyntheticDailyBarWriter.OHLCV:
result = self.bcolz_daily_bar_ctable
for column in OHLCV:
idx = 0
data = result[column][:]
multiplier = 1 if column == 'volume' else 1000
for asset_id in self.assets:
for date in self.dates_for_asset(asset_id):
self.assertEqual(
SyntheticDailyBarWriter.expected_value(
expected_daily_bar_value(
asset_id,
date,
column
@@ -134,7 +138,7 @@ class BcolzDailyBarTestCase(TestCase):
self.assertEqual(idx, len(data))
def test_write_day_and_id(self):
result = self.writer.write(self.dest, self.trading_days, self.assets)
result = self.bcolz_daily_bar_ctable
idx = 0
ids = result['id']
days = result['day']
@@ -145,7 +149,7 @@ class BcolzDailyBarTestCase(TestCase):
idx += 1
def test_write_attrs(self):
result = self.writer.write(self.dest, self.trading_days, self.assets)
result = self.bcolz_daily_bar_ctable
expected_first_row = {
'1': 0,
'2': 5, # Asset 1 has 5 trading days.
@@ -182,16 +186,19 @@ class BcolzDailyBarTestCase(TestCase):
)
def _check_read_results(self, columns, assets, start_date, end_date):
table = self.writer.write(self.dest, self.trading_days, self.assets)
reader = BcolzDailyBarReader(table)
results = reader.load_raw_arrays(columns, start_date, end_date, assets)
results = self.bcolz_daily_bar_reader.load_raw_arrays(
columns,
start_date,
end_date,
assets,
)
dates = self.trading_days_between(start_date, end_date)
for column, result in zip(columns, results):
assert_array_equal(
result,
self.writer.expected_values_2d(
expected_daily_bar_values_2d(
dates,
assets,
EQUITY_INFO,
column.name,
)
)
@@ -267,35 +274,34 @@ class BcolzDailyBarTestCase(TestCase):
)
def test_unadjusted_spot_price(self):
table = self.writer.write(self.dest, self.trading_days, self.assets)
reader = BcolzDailyBarReader(table)
reader = self.bcolz_daily_bar_reader
# At beginning
price = reader.spot_price(1, Timestamp('2015-06-01', tz='UTC'),
'close')
# Synthetic writes price for date.
self.assertEqual(135630.0, price)
self.assertEqual(108630.0, price)
# Middle
price = reader.spot_price(1, Timestamp('2015-06-02', tz='UTC'),
'close')
self.assertEqual(135631.0, price)
self.assertEqual(108631.0, price)
# End
price = reader.spot_price(1, Timestamp('2015-06-05', tz='UTC'),
'close')
self.assertEqual(135634.0, price)
self.assertEqual(108634.0, price)
# Another sid at beginning.
price = reader.spot_price(2, Timestamp('2015-06-22', tz='UTC'),
'close')
self.assertEqual(235651.0, price)
self.assertEqual(208651.0, price)
# Ensure that volume does not have float adjustment applied.
volume = reader.spot_price(1, Timestamp('2015-06-02', tz='UTC'),
'volume')
self.assertEqual(145631, volume)
self.assertEqual(109631, volume)
def test_unadjusted_spot_price_no_data(self):
table = self.writer.write(self.dest, self.trading_days, self.assets)
table = self.bcolz_daily_bar_ctable
reader = BcolzDailyBarReader(table)
# before
with self.assertRaises(NoDataOnDate):
@@ -306,18 +312,21 @@ class BcolzDailyBarTestCase(TestCase):
reader.spot_price(4, Timestamp('2015-06-16', tz='UTC'), 'close')
def test_unadjusted_spot_price_empty_value(self):
table = self.writer.write(self.dest, self.trading_days, self.assets)
reader = BcolzDailyBarReader(table)
reader = self.bcolz_daily_bar_reader
# A sid, day and corresponding index into which to overwrite a zero.
zero_sid = 1
zero_day = Timestamp('2015-06-02', tz='UTC')
zero_ix = reader.sid_day_index(zero_sid, zero_day)
# Write a zero into the synthetic pricing data at the day and sid,
# so that a read should now return -1.
# This a little hacky, in lieu of changing the synthetic data set.
reader._spot_col('close')[zero_ix] = 0
old = reader._spot_col('close')[zero_ix]
try:
# Write a zero into the synthetic pricing data at the day and sid,
# so that a read should now return -1.
# This a little hacky, in lieu of changing the synthetic data set.
reader._spot_col('close')[zero_ix] = 0
close = reader.spot_price(zero_sid, zero_day, 'close')
self.assertEqual(-1, close)
close = reader.spot_price(zero_sid, zero_day, 'close')
self.assertEqual(-1, close)
finally:
reader._spot_col('close')[zero_ix] = old
+78 -130
View File
@@ -8,133 +8,95 @@
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
'''
Unit tests for finance.slippage
"""
'''
import datetime
import pytz
from unittest import TestCase
from nose_parameterized import parameterized
import numpy as np
import pandas as pd
from pandas.tslib import normalize_date
from testfixtures import TempDirectory
from zipline.finance.slippage import VolumeShareSlippage
from zipline.finance.trading import TradingEnvironment, SimulationParameters
from zipline.protocol import DATASOURCE_TYPE
from zipline.finance.blotter import Order
from zipline.data.minute_bars import BcolzMinuteBarReader
from zipline.data.data_portal import DataPortal
from zipline.protocol import BarData
from zipline.testing.core import write_bcolz_minute_data
from zipline.testing import tmp_bcolz_minute_bar_reader
from zipline.testing.fixtures import (
WithDataPortal,
WithSimParams,
ZiplineTestCase,
)
class SlippageTestCase(TestCase):
class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-05 14:31', tz='utc')
END_DATE = pd.Timestamp('2006-01-05 14:36', tz='utc')
SIM_PARAMS_CAPITAL_BASE = 1.0e5
SIM_PARAMS_DATA_FREQUENCY = 'minute'
SIM_PARAMS_EMISSION_RATE = 'daily'
ASSET_FINDER_EQUITY_SIDS = (133,)
ASSET_FINDER_EQUITY_START_DATE = pd.Timestamp('2006-01-05', tz='utc')
ASSET_FINDER_EQUITY_END_DATE = pd.Timestamp('2006-01-07', tz='utc')
minutes = pd.DatetimeIndex(
start=START_DATE,
end=END_DATE - pd.Timedelta('1 minute'),
freq='1min'
)
@classmethod
def setUpClass(cls):
cls.tempdir = TempDirectory()
cls.env = TradingEnvironment()
cls.sim_params = SimulationParameters(
period_start=pd.Timestamp("2006-01-05 14:31", tz="utc"),
period_end=pd.Timestamp("2006-01-05 14:36", tz="utc"),
capital_base=1.0e5,
data_frequency="minute",
emission_rate='daily',
env=cls.env
)
cls.sids = [133]
cls.minutes = pd.DatetimeIndex(
start=pd.Timestamp("2006-01-05 14:31", tz="utc"),
end=pd.Timestamp("2006-01-05 14:35", tz="utc"),
freq="1min"
)
assets = {
133: pd.DataFrame({
"open": np.array([3.0, 3.0, 3.5, 4.0, 3.5]),
"high": np.array([3.15, 3.15, 3.15, 3.15, 3.15]),
"low": np.array([2.85, 2.85, 2.85, 2.85, 2.85]),
"close": np.array([3.0, 3.5, 4.0, 3.5, 3.0]),
"volume": [2000, 2000, 2000, 2000, 2000],
"dt": cls.minutes
}).set_index("dt")
def make_minute_bar_data(cls):
return {
133: pd.DataFrame(
{
'open': [3.0, 3.0, 3.5, 4.0, 3.5],
'high': [3.15, 3.15, 3.15, 3.15, 3.15],
'low': [2.85, 2.85, 2.85, 2.85, 2.85],
'close': [3.0, 3.5, 4.0, 3.5, 3.0],
'volume': [2000, 2000, 2000, 2000, 2000],
},
index=cls.minutes,
),
}
write_bcolz_minute_data(
cls.env,
pd.date_range(
start=normalize_date(cls.minutes[0]),
end=normalize_date(cls.minutes[-1])
),
cls.tempdir.path,
assets
)
cls.env.write_data(equities_data={
133: {
"start_date": pd.Timestamp("2006-01-05", tz='utc'),
"end_date": pd.Timestamp("2006-01-07", tz='utc')
}
})
@classmethod
def init_class_fixtures(cls):
super(SlippageTestCase, cls).init_class_fixtures()
cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133)
cls.data_portal = DataPortal(
cls.env,
equity_minute_reader=BcolzMinuteBarReader(cls.tempdir.path),
)
@classmethod
def tearDownClass(cls):
cls.tempdir.cleanup()
del cls.env
def test_volume_share_slippage(self):
tempdir = TempDirectory()
try:
assets = {
133: pd.DataFrame({
"open": [3.00],
"high": [3.15],
"low": [2.85],
"close": [3.00],
"volume": [200],
"dt": [self.minutes[0]]
}).set_index("dt")
}
write_bcolz_minute_data(
self.env,
pd.date_range(
start=normalize_date(self.minutes[0]),
end=normalize_date(self.minutes[-1])
),
tempdir.path,
assets
)
equity_minute_reader = BcolzMinuteBarReader(tempdir.path)
assets = {
133: pd.DataFrame(
{
'open': [3.00],
'high': [3.15],
'low': [2.85],
'close': [3.00],
'volume': [200],
},
index=[self.minutes[0]],
),
}
days = pd.date_range(
start=normalize_date(self.minutes[0]),
end=normalize_date(self.minutes[-1])
)
with tmp_bcolz_minute_bar_reader(self.env, days, assets) as reader:
data_portal = DataPortal(
self.env,
equity_minute_reader=equity_minute_reader,
equity_minute_reader=reader,
)
slippage_model = VolumeShareSlippage()
@@ -201,9 +163,6 @@ class SlippageTestCase(TestCase):
self.assertEquals(len(orders_txns), 0)
finally:
tempdir.cleanup()
def test_orders_limit(self):
slippage_model = VolumeShareSlippage()
slippage_model.data_portal = self.data_portal
@@ -502,39 +461,30 @@ class SlippageTestCase(TestCase):
for name, case in STOP_ORDER_CASES.items()
])
def test_orders_stop(self, name, order_data, event_data, expected):
tempdir = TempDirectory()
try:
data = order_data
data['sid'] = self.ASSET133
order = Order(**data)
assets = {
133: pd.DataFrame({
"open": [event_data["open"]],
"high": [event_data["high"]],
"low": [event_data["low"]],
"close": [event_data["close"]],
"volume": [event_data["volume"]],
"dt": [pd.Timestamp('2006-01-05 14:31', tz='UTC')]
}).set_index("dt")
}
write_bcolz_minute_data(
self.env,
pd.date_range(
start=normalize_date(self.minutes[0]),
end=normalize_date(self.minutes[-1])
),
tempdir.path,
assets
)
equity_minute_reader = BcolzMinuteBarReader(tempdir.path)
data = order_data
data['sid'] = self.ASSET133
order = Order(**data)
assets = {
133: pd.DataFrame(
{
'open': [event_data['open']],
'high': [event_data['high']],
'low': [event_data['low']],
'close': [event_data['close']],
'volume': [event_data['volume']],
},
index=[pd.Timestamp('2006-01-05 14:31', tz='UTC')],
),
}
days = pd.date_range(
start=normalize_date(self.minutes[0]),
end=normalize_date(self.minutes[-1])
)
with tmp_bcolz_minute_bar_reader(self.env, days, assets) as reader:
data_portal = DataPortal(
self.env,
equity_minute_reader=equity_minute_reader,
equity_minute_reader=reader,
)
slippage_model = VolumeShareSlippage()
@@ -559,8 +509,6 @@ class SlippageTestCase(TestCase):
for key, value in expected['transaction'].items():
self.assertEquals(value, txn[key])
finally:
tempdir.cleanup()
def test_orders_stop_limit(self):
slippage_model = VolumeShareSlippage()
+2 -2
View File
@@ -9,13 +9,13 @@ from numpy import arange, prod
from pandas import date_range, Int64Index, DataFrame
from six import iteritems
from zipline.pipeline import TermGraph
from zipline.assets.synthetic import make_simple_equity_info
from zipline.pipeline.engine import SimplePipelineEngine
from zipline.pipeline import TermGraph
from zipline.pipeline.term import AssetExists
from zipline.testing import (
check_arrays,
ExplodingObject,
make_simple_equity_info,
tmp_asset_finder,
)
+2 -4
View File
@@ -20,6 +20,7 @@ from pandas.util.testing import assert_frame_equal
from toolz import keymap, valmap, concatv
from toolz.curried import operator as op
from zipline.assets.synthetic import make_simple_equity_info
from zipline.pipeline import Pipeline, CustomFactor
from zipline.pipeline.data import DataSet, BoundColumn
from zipline.pipeline.engine import SimplePipelineEngine
@@ -38,10 +39,7 @@ from zipline.utils.numpy_utils import (
int64_dtype,
repeat_last_axis,
)
from zipline.testing import (
tmp_asset_finder,
make_simple_equity_info,
)
from zipline.testing import tmp_asset_finder
nameof = op.attrgetter('name')
dtypeof = op.attrgetter('dtype')
-1
View File
@@ -26,7 +26,6 @@ from zipline.pipeline.loaders.utils import (
get_values_for_date_ranges,
zip_with_dates
)
from zipline.testing.fixtures import (
WithPipelineEventDataLoader,
ZiplineTestCase
+83 -118
View File
@@ -3,7 +3,6 @@ Tests for SimplePipelineEngine
"""
from __future__ import division
from collections import OrderedDict
from unittest import TestCase
from itertools import product
from nose_parameterized import parameterized
@@ -35,24 +34,22 @@ from pandas import (
from pandas.compat.chainmap import ChainMap
from pandas.util.testing import assert_frame_equal
from six import iteritems, itervalues
from testfixtures import TempDirectory
from toolz import merge
from zipline.data.us_equity_pricing import BcolzDailyBarReader
from zipline.finance.trading import TradingEnvironment
from zipline.assets.synthetic import make_rotating_equity_info
from zipline.lib.adjustment import MULTIPLY
from zipline.pipeline.loaders.synthetic import (
PrecomputedLoader,
NullAdjustmentReader,
SyntheticDailyBarWriter,
)
from zipline.pipeline.loaders.synthetic import PrecomputedLoader
from zipline.pipeline import Pipeline
from zipline.pipeline.data import USEquityPricing, DataSet, Column
from zipline.pipeline.loaders.frame import DataFrameLoader
from zipline.pipeline.loaders.equity_pricing_loader import (
USEquityPricingLoader,
)
from zipline.pipeline.loaders.synthetic import (
make_daily_bar_data,
expected_daily_bar_values_2d,
)
from zipline.pipeline.engine import SimplePipelineEngine
from zipline.pipeline.loaders.frame import DataFrameLoader
from zipline.pipeline import CustomFactor
from zipline.pipeline.factors import (
AverageDollarVolume,
@@ -64,11 +61,14 @@ from zipline.pipeline.factors import (
SimpleMovingAverage,
)
from zipline.testing import (
make_rotating_equity_info,
make_simple_equity_info,
product_upper_triangle,
check_arrays,
)
from zipline.testing.fixtures import (
WithAdjustmentReader,
WithTradingEnvironment,
ZiplineTestCase,
)
from zipline.utils.memoize import lazyval
@@ -163,10 +163,15 @@ class RollingSumSum(CustomFactor):
out[:] = sum(inputs).sum(axis=0)
class ConstantInputTestCase(TestCase):
class ConstantInputTestCase(WithTradingEnvironment, ZiplineTestCase):
asset_ids = ASSET_FINDER_EQUITY_SIDS = 1, 2, 3, 4
START_DATE = Timestamp('2014-01-01', tz='utc')
END_DATE = Timestamp('2014-03-01', tz='utc')
def setUp(self):
self.constants = {
@classmethod
def init_class_fixtures(cls):
super(ConstantInputTestCase, cls).init_class_fixtures()
cls.constants = {
# Every day, assume every stock starts at 2, goes down to 1,
# goes up to 4, and finishes at 3.
USEquityPricing.low: 1,
@@ -174,23 +179,18 @@ class ConstantInputTestCase(TestCase):
USEquityPricing.close: 3,
USEquityPricing.high: 4,
}
self.asset_ids = [1, 2, 3, 4]
self.dates = date_range('2014-01', '2014-03', freq='D', tz='UTC')
self.loader = PrecomputedLoader(
constants=self.constants,
dates=self.dates,
sids=self.asset_ids,
cls.dates = date_range(
cls.START_DATE,
cls.END_DATE,
freq='D',
tz='UTC',
)
self.asset_info = make_simple_equity_info(
self.asset_ids,
start_date=self.dates[0],
end_date=self.dates[-1],
cls.loader = PrecomputedLoader(
constants=cls.constants,
dates=cls.dates,
sids=cls.asset_ids,
)
environment = TradingEnvironment()
environment.write_data(equities_df=self.asset_info)
self.asset_finder = environment.asset_finder
self.assets = self.asset_finder.retrieve_all(self.asset_ids)
cls.assets = cls.asset_finder.retrieve_all(cls.asset_ids)
def test_bad_dates(self):
loader = self.loader
@@ -608,35 +608,22 @@ class ConstantInputTestCase(TestCase):
Loader2DataSet.col2)})
class FrameInputTestCase(TestCase):
class FrameInputTestCase(WithTradingEnvironment, ZiplineTestCase):
asset_ids = ASSET_FINDER_EQUITY_SIDS = 1, 2, 3
start = START_DATE = Timestamp('2015-01-01', tz='utc')
end = END_DATE = Timestamp('2015-01-31', tz='utc')
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
day = cls.env.trading_day
cls.asset_ids = [1, 2, 3]
def init_class_fixtures(cls):
super(FrameInputTestCase, cls).init_class_fixtures()
cls.dates = date_range(
'2015-01-01',
'2015-01-31',
freq=day,
cls.start,
cls.end,
freq=cls.env.trading_day,
tz='UTC',
)
asset_info = make_simple_equity_info(
cls.asset_ids,
start_date=cls.dates[0],
end_date=cls.dates[-1],
)
cls.env.write_data(equities_df=asset_info)
cls.asset_finder = cls.env.asset_finder
cls.assets = cls.asset_finder.retrieve_all(cls.asset_ids)
@classmethod
def tearDownClass(cls):
del cls.env
del cls.asset_finder
@lazyval
def base_mask(self):
return self.make_frame(True)
@@ -725,54 +712,39 @@ class FrameInputTestCase(TestCase):
assert_frame_equal(high_results, high_base.iloc[iloc_bounds])
class SyntheticBcolzTestCase(TestCase):
class SyntheticBcolzTestCase(WithAdjustmentReader,
ZiplineTestCase):
first_asset_start = Timestamp('2015-04-01', tz='UTC')
START_DATE = Timestamp('2015-01-01', tz='utc')
END_DATE = Timestamp('2015-08-01', tz='utc')
@classmethod
def setUpClass(cls):
cls.first_asset_start = Timestamp('2015-04-01', tz='UTC')
cls.env = TradingEnvironment()
cls.trading_day = day = cls.env.trading_day
cls.calendar = date_range('2015', '2015-08', tz='UTC', freq=day)
cls.asset_info = make_rotating_equity_info(
def make_equity_info(cls):
cls.equity_info = ret = make_rotating_equity_info(
num_assets=6,
first_start=cls.first_asset_start,
frequency=day,
frequency=cls.TRADING_ENV_TRADING_CALENDAR.trading_day,
periods_between_starts=4,
asset_lifetime=8,
)
cls.last_asset_end = cls.asset_info['end_date'].max()
cls.all_asset_ids = cls.asset_info.index
cls.env.write_data(equities_df=cls.asset_info)
cls.finder = cls.env.asset_finder
cls.temp_dir = TempDirectory()
cls.temp_dir.create()
try:
cls.writer = SyntheticDailyBarWriter(
asset_info=cls.asset_info[['start_date', 'end_date']],
calendar=cls.calendar,
)
table = cls.writer.write(
cls.temp_dir.getpath('testdata.bcolz'),
cls.calendar,
cls.all_asset_ids,
)
cls.pipeline_loader = USEquityPricingLoader(
BcolzDailyBarReader(table),
NullAdjustmentReader(),
)
except:
cls.temp_dir.cleanup()
raise
return ret
@classmethod
def tearDownClass(cls):
del cls.env
cls.temp_dir.cleanup()
def make_daily_bar_data(cls):
return make_daily_bar_data(
cls.equity_info,
cls.bcolz_daily_bar_days,
)
@classmethod
def init_class_fixtures(cls):
super(SyntheticBcolzTestCase, cls).init_class_fixtures()
cls.all_asset_ids = cls.asset_finder.sids
cls.last_asset_end = cls.equity_info['end_date'].max()
cls.pipeline_loader = USEquityPricingLoader(
cls.bcolz_daily_bar_reader,
cls.adjustment_reader,
)
def write_nans(self, df):
"""
@@ -807,14 +779,14 @@ class SyntheticBcolzTestCase(TestCase):
engine = SimplePipelineEngine(
lambda column: self.pipeline_loader,
self.env.trading_days,
self.finder,
self.asset_finder,
)
window_length = 5
asset_ids = self.all_asset_ids
dates = date_range(
self.first_asset_start + self.trading_day,
self.first_asset_start + self.env.trading_day,
self.last_asset_end,
freq=self.trading_day,
freq=self.env.trading_day,
)
dates_to_test = dates[window_length:]
@@ -833,8 +805,10 @@ class SyntheticBcolzTestCase(TestCase):
# computed results to be computed using values anchored on the
# **previous** day's data.
expected_raw = rolling_mean(
self.writer.expected_values_2d(
dates - self.trading_day, asset_ids, 'close',
expected_daily_bar_values_2d(
dates - self.env.trading_day,
self.equity_info,
'close',
),
window_length,
min_periods=1,
@@ -844,7 +818,7 @@ class SyntheticBcolzTestCase(TestCase):
# Truncate off the extra rows needed to compute the SMAs.
expected_raw[window_length:],
index=dates_to_test, # dates_to_test is dates[window_length:]
columns=self.finder.retrieve_all(asset_ids),
columns=self.asset_finder.retrieve_all(asset_ids),
)
self.write_nans(expected)
result = results['sma'].unstack()
@@ -859,14 +833,14 @@ class SyntheticBcolzTestCase(TestCase):
engine = SimplePipelineEngine(
lambda column: self.pipeline_loader,
self.env.trading_days,
self.finder,
self.asset_finder,
)
window_length = 5
asset_ids = self.all_asset_ids
dates = date_range(
self.first_asset_start + self.trading_day,
self.first_asset_start + self.env.trading_day,
self.last_asset_end,
freq=self.trading_day,
freq=self.env.trading_day,
)
dates_to_test = dates[window_length:]
@@ -886,7 +860,7 @@ class SyntheticBcolzTestCase(TestCase):
expected = DataFrame(
data=zeros((len(dates_to_test), len(asset_ids)), dtype=float),
index=dates_to_test,
columns=self.finder.retrieve_all(asset_ids),
columns=self.asset_finder.retrieve_all(asset_ids),
)
self.write_nans(expected)
result = results['drawdown'].unstack()
@@ -894,27 +868,23 @@ class SyntheticBcolzTestCase(TestCase):
assert_frame_equal(expected, result)
class ParameterizedFactorTestCase(TestCase):
class ParameterizedFactorTestCase(WithTradingEnvironment, ZiplineTestCase):
sids = ASSET_FINDER_EQUITY_SIDS = Int64Index([1, 2, 3])
START_DATE = Timestamp('2015-01-31', tz='UTC')
END_DATE = Timestamp('2015-03-01', tz='UTC')
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
def init_class_fixtures(cls):
super(ParameterizedFactorTestCase, cls).init_class_fixtures()
day = cls.env.trading_day
cls.sids = sids = Int64Index([1, 2, 3])
cls.dates = dates = date_range(
'2015-02-01',
'2015-02-28',
freq=day,
tz='UTC',
)
asset_info = make_simple_equity_info(
cls.sids,
start_date=Timestamp('2015-01-31', tz='UTC'),
end_date=Timestamp('2015-03-01', tz='UTC'),
)
cls.env.write_data(equities_df=asset_info)
cls.asset_finder = cls.env.asset_finder
sids = cls.sids
cls.raw_data = DataFrame(
data=arange(len(dates) * len(sids), dtype=float).reshape(
@@ -939,11 +909,6 @@ class ParameterizedFactorTestCase(TestCase):
cls.asset_finder,
)
@classmethod
def tearDownClass(cls):
del cls.env
del cls.asset_finder
def expected_ewma(self, window_length, decay_rate):
alpha = 1 - decay_rate
span = (2 / alpha) - 1
+100 -141
View File
@@ -1,8 +1,6 @@
"""
Tests for Algorithms using the Pipeline API.
"""
import os
from unittest import TestCase
from os.path import (
dirname,
join,
@@ -19,6 +17,7 @@ from numpy import (
uint32,
)
from numpy.testing import assert_almost_equal
import pandas as pd
from pandas import (
concat,
DataFrame,
@@ -28,7 +27,6 @@ from pandas import (
Timestamp,
)
from six import iteritems, itervalues
from testfixtures import TempDirectory
from zipline.algorithm import TradingAlgorithm
from zipline.api import (
@@ -36,19 +34,11 @@ from zipline.api import (
pipeline_output,
get_datetime,
)
from zipline.data.data_portal import DataPortal
from zipline.errors import (
AttachPipelineAfterInitialize,
PipelineOutputDuringInitialize,
NoSuchPipeline,
)
from zipline.data.us_equity_pricing import (
BcolzDailyBarReader,
DailyBarWriterFromCSVs,
SQLiteAdjustmentWriter,
SQLiteAdjustmentReader,
)
from zipline.finance import trading
from zipline.lib.adjustment import MULTIPLY
from zipline.pipeline import Pipeline
from zipline.pipeline.factors import VWAP
@@ -58,15 +48,19 @@ from zipline.pipeline.loaders.equity_pricing_loader import (
USEquityPricingLoader,
)
from zipline.testing import (
make_simple_equity_info,
str_to_seconds
)
from zipline.testing.core import DailyBarWriterFromDataFrames, \
create_empty_splits_mergers_frame, FakeDataPortal
from zipline.utils.tradingcalendar import (
trading_day,
trading_days,
from zipline.testing import (
create_empty_splits_mergers_frame,
FakeDataPortal,
)
from zipline.testing.fixtures import (
WithAdjustmentReader,
WithBcolzDailyBarReaderFromCSVs,
WithDataPortal,
ZiplineTestCase,
)
from zipline.utils.tradingcalendar import trading_day
TEST_RESOURCE_PATH = join(
@@ -89,101 +83,83 @@ def rolling_vwap(df, length):
return Series(out, index=df.index)
class ClosesOnly(TestCase):
class ClosesOnly(WithDataPortal, ZiplineTestCase):
sids = 1, 2, 3
START_DATE = pd.Timestamp('2014-01-01', tz='utc')
END_DATE = pd.Timestamp('2014-02-01', tz='utc')
dates = date_range(START_DATE, END_DATE, freq=trading_day, tz='utc')
@classmethod
def setUpClass(cls):
cls.tempdir = TempDirectory()
@classmethod
def tearDownClass(cls):
cls.tempdir.cleanup()
def setUp(self):
self.env = env = trading.TradingEnvironment()
self.dates = date_range(
'2014-01-01', '2014-02-01', freq=trading_day, tz='UTC'
)
asset_info = DataFrame.from_records([
def make_equity_info(cls):
cls.equity_info = ret = DataFrame.from_records([
{
'sid': 1,
'symbol': 'A',
'start_date': self.dates[10],
'end_date': self.dates[13],
'start_date': cls.dates[10],
'end_date': cls.dates[13],
'exchange': 'TEST',
},
{
'sid': 2,
'symbol': 'B',
'start_date': self.dates[11],
'end_date': self.dates[14],
'start_date': cls.dates[11],
'end_date': cls.dates[14],
'exchange': 'TEST',
},
{
'sid': 3,
'symbol': 'C',
'start_date': self.dates[12],
'end_date': self.dates[15],
'start_date': cls.dates[12],
'end_date': cls.dates[15],
'exchange': 'TEST',
},
])
self.first_asset_start = min(asset_info.start_date)
self.last_asset_end = max(asset_info.end_date)
env.write_data(equities_df=asset_info)
self.asset_finder = finder = env.asset_finder
return ret
sids = (1, 2, 3)
self.assets = finder.retrieve_all(sids)
# View of the baseline data.
self.closes = DataFrame(
{sid: arange(1, len(self.dates) + 1) * sid for sid in sids},
index=self.dates,
@classmethod
def make_daily_bar_data(cls):
cls.closes = DataFrame(
{sid: arange(1, len(cls.dates) + 1) * sid for sid in cls.sids},
index=cls.dates,
dtype=float,
)
for sid in cls.sids:
yield sid, DataFrame(
{
'open': cls.closes[sid].values,
'high': cls.closes[sid].values,
'low': cls.closes[sid].values,
'close': cls.closes[sid].values,
'volume': cls.closes[sid].values,
},
index=cls.dates,
)
# Create a data portal holding the data in self.closes
data = {}
for sid in sids:
data[sid] = DataFrame({
"open": self.closes[sid].values,
"high": self.closes[sid].values,
"low": self.closes[sid].values,
"close": self.closes[sid].values,
"volume": self.closes[sid].values,
"day": [day.value for day in self.dates]
})
path = os.path.join(self.tempdir.path, "testdaily.bcolz")
DailyBarWriterFromDataFrames(data).write(
path,
self.dates,
data
)
daily_bar_reader = BcolzDailyBarReader(path)
self.data_portal = DataPortal(
self.env,
equity_daily_reader=daily_bar_reader,
)
@classmethod
def init_class_fixtures(cls):
super(ClosesOnly, cls).init_class_fixtures()
cls.first_asset_start = min(cls.equity_info.start_date)
cls.last_asset_end = max(cls.equity_info.end_date)
cls.assets = cls.asset_finder.retrieve_all(cls.sids)
# Add a split for 'A' on its second date.
self.split_asset = self.assets[0]
self.split_date = self.split_asset.start_date + trading_day
self.split_ratio = 0.5
self.adjustments = DataFrame.from_records([
cls.split_asset = cls.assets[0]
cls.split_date = cls.split_asset.start_date + trading_day
cls.split_ratio = 0.5
cls.adjustments = DataFrame.from_records([
{
'sid': self.split_asset.sid,
'value': self.split_ratio,
'sid': cls.split_asset.sid,
'value': cls.split_ratio,
'kind': MULTIPLY,
'start_date': Timestamp('NaT'),
'end_date': self.split_date,
'apply_date': self.split_date,
'end_date': cls.split_date,
'apply_date': cls.split_date,
}
])
def init_instance_fixtures(self):
super(ClosesOnly, self).init_instance_fixtures()
# View of the data on/after the split.
self.adj_closes = adj_closes = self.closes.copy()
adj_closes.ix[:self.split_date, self.split_asset] *= self.split_ratio
@@ -361,87 +337,70 @@ class MockDailyBarSpotReader(object):
return 100.0
class PipelineAlgorithmTestCase(TestCase):
class PipelineAlgorithmTestCase(WithBcolzDailyBarReaderFromCSVs,
WithAdjustmentReader,
ZiplineTestCase):
AAPL = 1
MSFT = 2
BRK_A = 3
assets = ASSET_FINDER_EQUITY_SIDS = AAPL, MSFT, BRK_A
ASSET_FINDER_EQUITY_SYMBOLS = 'AAPL', 'MSFT', 'BRK_A'
START_DATE = Timestamp('2014')
END_DATE = Timestamp('2015')
BCOLZ_DAILY_BAR_USE_FULL_CALENDAR = True
@classmethod
def setUpClass(cls):
cls.AAPL = 1
cls.MSFT = 2
cls.BRK_A = 3
cls.assets = [cls.AAPL, cls.MSFT, cls.BRK_A]
asset_info = make_simple_equity_info(
cls.assets,
Timestamp('2014'),
Timestamp('2015'),
['AAPL', 'MSFT', 'BRK_A'],
)
cls.env = trading.TradingEnvironment()
cls.env.write_data(equities_df=asset_info)
cls.tempdir = tempdir = TempDirectory()
tempdir.create()
try:
cls.raw_data, bar_reader = cls.create_bar_reader(tempdir)
adj_reader = cls.create_adjustment_reader(tempdir)
cls.pipeline_loader = USEquityPricingLoader(
bar_reader, adj_reader
)
except:
cls.tempdir.cleanup()
raise
cls.dates = cls.raw_data[cls.AAPL].index.tz_localize('UTC')
cls.AAPL_split_date = Timestamp("2014-06-09", tz='UTC')
@classmethod
def tearDownClass(cls):
del cls.pipeline_loader
del cls.env
cls.tempdir.cleanup()
@classmethod
def create_bar_reader(cls, tempdir):
def make_daily_bar_data(cls):
resources = {
cls.AAPL: join(TEST_RESOURCE_PATH, 'AAPL.csv'),
cls.MSFT: join(TEST_RESOURCE_PATH, 'MSFT.csv'),
cls.BRK_A: join(TEST_RESOURCE_PATH, 'BRK-A.csv'),
}
raw_data = {
cls.raw_data = raw_data = {
asset: read_csv(path, parse_dates=['day']).set_index('day')
for asset, path in iteritems(resources)
for asset, path in resources.items()
}
# Add 'price' column as an alias because all kinds of stuff in zipline
# depends on it being present. :/
for frame in raw_data.values():
frame['price'] = frame['close']
writer = DailyBarWriterFromCSVs(resources)
data_path = tempdir.getpath('testdata.bcolz')
table = writer.write(data_path, trading_days, cls.assets)
return raw_data, BcolzDailyBarReader(table)
return resources
@classmethod
def create_adjustment_reader(cls, tempdir):
dbpath = tempdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, cls.env.trading_days,
MockDailyBarSpotReader())
splits = DataFrame.from_records([
def make_splits_data(cls):
return DataFrame.from_records([
{
'effective_date': str_to_seconds('2014-06-09'),
'ratio': (1 / 7.0),
'sid': cls.AAPL,
}
])
mergers = create_empty_splits_mergers_frame()
dividends = DataFrame({
'sid': array([], dtype=uint32),
'amount': array([], dtype=float64),
'record_date': array([], dtype='datetime64[ns]'),
'ex_date': array([], dtype='datetime64[ns]'),
'declared_date': array([], dtype='datetime64[ns]'),
'pay_date': array([], dtype='datetime64[ns]'),
})
writer.write(splits, mergers, dividends)
return SQLiteAdjustmentReader(dbpath)
@classmethod
def make_mergers_data(cls):
return create_empty_splits_mergers_frame()
@classmethod
def make_dividends_data(cls):
return pd.DataFrame(array([], dtype=[
('sid', uint32),
('amount', float64),
('record_date', 'datetime64[ns]'),
('ex_date', 'datetime64[ns]'),
('declared_date', 'datetime64[ns]'),
('pay_date', 'datetime64[ns]'),
]))
@classmethod
def init_class_fixtures(cls):
super(PipelineAlgorithmTestCase, cls).init_class_fixtures()
cls.pipeline_loader = USEquityPricingLoader(
cls.bcolz_daily_bar_reader,
cls.adjustment_reader,
)
cls.dates = cls.raw_data[cls.AAPL].index.tz_localize('UTC')
cls.AAPL_split_date = Timestamp("2014-06-09", tz='UTC')
def compute_expected_vwaps(self, window_lengths):
AAPL, MSFT, BRK_A = self.AAPL, self.MSFT, self.BRK_A
+54 -58
View File
@@ -15,8 +15,6 @@
"""
Tests for USEquityPricingLoader and related classes.
"""
from unittest import TestCase
from numpy import (
arange,
datetime64,
@@ -34,29 +32,28 @@ from pandas import (
Int64Index,
Timestamp,
)
from testfixtures import TempDirectory
from toolz.curried.operator import getitem
from zipline.lib.adjustment import Float64Multiply
from zipline.pipeline.loaders.synthetic import (
NullAdjustmentReader,
SyntheticDailyBarWriter,
)
from zipline.data.us_equity_pricing import (
BcolzDailyBarReader,
SQLiteAdjustmentReader,
SQLiteAdjustmentWriter,
make_daily_bar_data,
expected_daily_bar_values_2d,
)
from zipline.pipeline.loaders.equity_pricing_loader import (
USEquityPricingLoader,
)
from zipline.errors import WindowLengthTooLong
from zipline.finance.trading import TradingEnvironment
from zipline.pipeline.data import USEquityPricing
from zipline.testing import (
seconds_to_timestamp,
str_to_seconds,
MockDailyBarReader,
)
from zipline.testing.fixtures import (
WithAdjustmentReader,
ZiplineTestCase,
)
# Test calendar ranges over the month of June 2015
@@ -257,41 +254,44 @@ DIVIDENDS_EXPECTED = DataFrame(
)
class MockDailyBarSpotReader(object):
"""
A BcolzDailyBarReader which returns a constant value for spot price.
"""
def spot_price(self, sid, day, column):
return 100.0
class USEquityPricingLoaderTestCase(TestCase):
class USEquityPricingLoaderTestCase(WithAdjustmentReader,
ZiplineTestCase):
START_DATE = TEST_CALENDAR_START
END_DATE = TEST_CALENDAR_STOP
asset_ids = 1, 2, 3
@classmethod
def setUpClass(cls):
cls.test_data_dir = TempDirectory()
cls.db_path = cls.test_data_dir.getpath('adjustments.db')
all_days = TradingEnvironment().trading_days
cls.calendar_days = all_days[
all_days.slice_indexer(TEST_CALENDAR_START, TEST_CALENDAR_STOP)
]
daily_bar_reader = MockDailyBarSpotReader()
writer = SQLiteAdjustmentWriter(cls.db_path, cls.calendar_days,
daily_bar_reader)
writer.write(SPLITS, MERGERS, DIVIDENDS)
def make_equity_info(cls):
return EQUITY_INFO
@classmethod
def make_splits_data(cls):
return SPLITS
@classmethod
def make_mergers_data(cls):
return MERGERS
@classmethod
def make_dividends_data(cls):
return DIVIDENDS
@classmethod
def make_adjustment_writer_daily_bar_reader(cls):
return MockDailyBarReader()
@classmethod
def make_daily_bar_data(cls):
return make_daily_bar_data(
EQUITY_INFO,
cls.bcolz_daily_bar_days,
)
@classmethod
def init_class_fixtures(cls):
super(USEquityPricingLoaderTestCase, cls).init_class_fixtures()
cls.assets = TEST_QUERY_ASSETS
cls.asset_info = EQUITY_INFO
cls.bcolz_writer = SyntheticDailyBarWriter(
cls.asset_info,
cls.calendar_days,
)
cls.bcolz_path = cls.test_data_dir.getpath('equity_pricing.bcolz')
cls.bcolz_writer.write(cls.bcolz_path, cls.calendar_days, cls.assets)
@classmethod
def tearDownClass(cls):
cls.test_data_dir.cleanup()
def test_input_sanity(self):
# Ensure that the input data doesn't contain adjustments during periods
@@ -306,13 +306,13 @@ class USEquityPricingLoaderTestCase(TestCase):
self.assertLessEqual(eff_date, asset_end)
def calendar_days_between(self, start_date, end_date, shift=0):
slice_ = self.calendar_days.slice_indexer(start_date, end_date)
slice_ = self.bcolz_daily_bar_days.slice_indexer(start_date, end_date)
start = slice_.start + shift
stop = slice_.stop + shift
if start < 0:
raise KeyError(start_date, shift)
return self.calendar_days[start:stop]
return self.bcolz_daily_bar_days[start:stop]
def expected_adjustments(self, start_date, end_date):
price_adjustments = {}
@@ -357,14 +357,13 @@ class USEquityPricingLoaderTestCase(TestCase):
return price_adjustments, volume_adjustments
def test_load_adjustments_from_sqlite(self):
reader = SQLiteAdjustmentReader(self.db_path)
columns = [USEquityPricing.close, USEquityPricing.volume]
query_days = self.calendar_days_between(
TEST_QUERY_START,
TEST_QUERY_STOP,
)
adjustments = reader.load_adjustments(
adjustments = self.adjustment_reader.load_adjustments(
columns,
query_days,
self.assets,
@@ -417,9 +416,8 @@ class USEquityPricingLoaderTestCase(TestCase):
)
self.assertEqual(adjustments, [{}, {}])
baseline_reader = BcolzDailyBarReader(self.bcolz_path)
pricing_loader = USEquityPricingLoader(
baseline_reader,
self.bcolz_daily_bar_reader,
adjustment_reader,
)
@@ -431,14 +429,14 @@ class USEquityPricingLoaderTestCase(TestCase):
)
closes, volumes = map(getitem(results), columns)
expected_baseline_closes = self.bcolz_writer.expected_values_2d(
expected_baseline_closes = expected_daily_bar_values_2d(
shifted_query_days,
self.assets,
self.asset_info,
'close',
)
expected_baseline_volumes = self.bcolz_writer.expected_values_2d(
expected_baseline_volumes = expected_daily_bar_values_2d(
shifted_query_days,
self.assets,
self.asset_info,
'volume',
)
@@ -495,11 +493,9 @@ class USEquityPricingLoaderTestCase(TestCase):
shift=-1,
)
baseline_reader = BcolzDailyBarReader(self.bcolz_path)
adjustment_reader = SQLiteAdjustmentReader(self.db_path)
pricing_loader = USEquityPricingLoader(
baseline_reader,
adjustment_reader,
self.bcolz_daily_bar_reader,
self.adjustment_reader,
)
results = pricing_loader.load_adjusted_array(
@@ -510,14 +506,14 @@ class USEquityPricingLoaderTestCase(TestCase):
)
highs, volumes = map(getitem(results), columns)
expected_baseline_highs = self.bcolz_writer.expected_values_2d(
expected_baseline_highs = expected_daily_bar_values_2d(
shifted_query_days,
self.assets,
self.asset_info,
'high',
)
expected_baseline_volumes = self.bcolz_writer.expected_values_2d(
expected_baseline_volumes = expected_daily_bar_values_2d(
shifted_query_days,
self.assets,
self.asset_info,
'volume',
)
@@ -68,41 +68,6 @@ aapl,1/9/06 5:31AM, 1
aapl,1/9/06 11:30AM, 4
""".strip()
IBM_CSV_DATA = """
symbol,date,signal
ibm,1/1/06,1
ibm,2/1/06,0
ibm,3/1/06,0
ibm,4/1/06,0
ibm,5/1/06,1
ibm,6/1/06,1
ibm,7/1/06,1
ibm,8/1/06,1
ibm,9/1/06,0
ibm,10/1/06,1
ibm,11/1/06,1
ibm,12/1/06,5
ibm,1/1/07,1
ibm,2/1/07,0
ibm,3/1/07,1
ibm,4/1/07,0
ibm,5/1/07,1
""".strip()
ANNUAL_AAPL_CSV_DATA = """
symbol,date,signal
aapl,1/2/06,1
aapl,1/15/06,1
aapl,3/1/06,1
aapl,3/15/06,1
aapl,6/1/06,1
aapl,6/15/06,1
aapl,9/1/06,1
aapl,9/15/06,1
aapl,12/1/06,1
aapl,12/15/06,1
""".strip()
AAPL_IBM_CSV_DATA = """
symbol,date,signal
aapl,1/1/06,1
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,39 @@
"""
Script for rebuilding the samples for the Quandl tests.
"""
from __future__ import print_function
import pandas as pd
import requests
from zipline.data.quandl import format_wiki_url
from zipline.utils.test_utils import test_resource_path, write_compressed
def zipfile_path(symbol):
return test_resource_path('quandl_samples', symbol + '.csv.gz')
def main():
start_date = pd.Timestamp('2014')
end_date = pd.Timestamp('2015')
symbols = ['AAPL', 'MSFT', 'BRK_A', 'ZEN']
for sym in symbols:
url = format_wiki_url(
api_key=None,
symbol=sym,
start_date=start_date,
end_date=end_date,
)
print("Fetching from %s" % url)
response = requests.get(url)
response.raise_for_status()
path = zipfile_path(sym)
print("Writing compressed data to %s" % path)
write_compressed(path, response.content)
if __name__ == '__main__':
main()
+487 -815
View File
File diff suppressed because it is too large Load Diff
+61 -133
View File
@@ -1,21 +1,23 @@
import warnings
from unittest import TestCase
from mock import patch
import pandas as pd
import numpy as np
from testfixtures import TempDirectory
import pandas as pd
from zipline import TradingAlgorithm
from zipline.data.data_portal import DataPortal
from zipline.data.minute_bars import BcolzMinuteBarWriter, \
US_EQUITIES_MINUTES_PER_DAY, BcolzMinuteBarReader
from zipline.data.us_equity_pricing import BcolzDailyBarReader, \
SQLiteAdjustmentReader, SQLiteAdjustmentWriter
from zipline.finance.trading import TradingEnvironment, SimulationParameters
from zipline.finance.trading import SimulationParameters
from zipline.protocol import BarData
from zipline.testing.core import write_minute_data_for_asset, \
create_daily_df_for_asset, DailyBarWriterFromDataFrames, MockDailyBarReader
from zipline.testing import str_to_seconds
from zipline.testing import (
MockDailyBarReader,
create_daily_df_for_asset,
create_minute_df_for_asset,
str_to_seconds,
)
from zipline.testing.fixtures import (
WithDataPortal,
WithSimParams,
ZiplineTestCase,
)
from zipline.zipline_warnings import ZiplineDeprecationWarning
simple_algo = """
@@ -111,137 +113,63 @@ def handle_data(context, data):
"""
class TestAPIShim(TestCase):
class TestAPIShim(WithDataPortal, WithSimParams, ZiplineTestCase):
START_DATE = pd.Timestamp("2016-01-05", tz='UTC')
END_DATE = pd.Timestamp("2016-01-28", tz='UTC')
SIM_PARAMS_DATA_FREQUENCY = 'minute'
sids = ASSET_FINDER_EQUITY_SIDS = 1, 2, 3
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.tempdir = TempDirectory()
def make_minute_bar_data(cls):
return {
sid: create_minute_df_for_asset(
cls.env,
cls.SIM_PARAMS_START,
cls.SIM_PARAMS_END,
)
for sid in cls.sids
}
cls.trading_days = cls.env.days_in_range(
start=pd.Timestamp("2016-01-05", tz='UTC'),
end=pd.Timestamp("2016-01-28", tz='UTC')
)
@classmethod
def make_daily_bar_data(cls):
for sid in cls.sids:
yield sid, create_daily_df_for_asset(
cls.env,
cls.SIM_PARAMS_START,
cls.SIM_PARAMS_END,
)
equities_data = {}
for sid in [1, 2, 3]:
equities_data[sid] = {
"start_date": cls.trading_days[0],
"end_date": cls.env.next_trading_day(cls.trading_days[-1]),
"symbol": "ASSET{0}".format(sid),
@classmethod
def make_splits_data(cls):
return pd.DataFrame([
{
'effective_date': str_to_seconds('2016-01-06'),
'ratio': 0.5,
'sid': 3,
}
])
cls.env.write_data(equities_data=equities_data)
@classmethod
def make_adjustment_writer_daily_bar_reader(cls):
return MockDailyBarReader()
@classmethod
def init_class_fixtures(cls):
super(TestAPIShim, cls).init_class_fixtures()
cls.asset1 = cls.env.asset_finder.retrieve_asset(1)
cls.asset2 = cls.env.asset_finder.retrieve_asset(2)
cls.asset3 = cls.env.asset_finder.retrieve_asset(3)
market_opens = cls.env.open_and_closes.market_open.loc[
cls.trading_days]
market_closes = cls.env.open_and_closes.market_close.loc[
cls.trading_days]
minute_writer = BcolzMinuteBarWriter(
cls.trading_days[0],
cls.tempdir.path,
market_opens,
market_closes,
US_EQUITIES_MINUTES_PER_DAY
)
for sid in [1, 2, 3]:
write_minute_data_for_asset(
cls.env, minute_writer, cls.trading_days[0],
cls.trading_days[-1], sid
)
cls.adj_reader = cls.create_adjustments_reader()
cls.sim_params = SimulationParameters(
period_start=cls.trading_days[0],
period_end=cls.trading_days[-1],
data_frequency="minute",
env=cls.env
)
@classmethod
def tearDownClass(cls):
del cls.adj_reader
cls.tempdir.cleanup()
@classmethod
def build_daily_data(cls):
path = cls.tempdir.getpath("testdaily.bcolz")
dfs = {
1: create_daily_df_for_asset(cls.env, cls.trading_days[0],
cls.trading_days[-1]),
2: create_daily_df_for_asset(cls.env, cls.trading_days[0],
cls.trading_days[-1]),
3: create_daily_df_for_asset(cls.env, cls.trading_days[0],
cls.trading_days[-1])
}
daily_writer = DailyBarWriterFromDataFrames(dfs)
daily_writer.write(path, cls.trading_days, dfs)
return BcolzDailyBarReader(path)
@classmethod
def create_adjustments_reader(cls):
path = cls.tempdir.getpath("test_adjustments.db")
adj_writer = SQLiteAdjustmentWriter(
path,
cls.env.trading_days,
MockDailyBarReader()
)
splits = pd.DataFrame([
{
'effective_date': str_to_seconds("2016-01-06"),
'ratio': 0.5,
'sid': cls.asset3.sid
}
])
# Mergers and Dividends are not tested, but we need to have these
# anyway
mergers = pd.DataFrame({}, columns=['effective_date', 'ratio', 'sid'])
mergers.effective_date = mergers.effective_date.astype(np.int64)
mergers.ratio = mergers.ratio.astype(np.float64)
mergers.sid = mergers.sid.astype(np.int64)
dividends = pd.DataFrame({}, columns=['ex_date', 'record_date',
'declared_date', 'pay_date',
'amount', 'sid'])
dividends.amount = dividends.amount.astype(np.float64)
dividends.sid = dividends.sid.astype(np.int64)
adj_writer.write(splits, mergers, dividends)
return SQLiteAdjustmentReader(path)
def setUp(self):
self.data_portal = DataPortal(
self.env,
equity_minute_reader=BcolzMinuteBarReader(self.tempdir.path),
equity_daily_reader=self.build_daily_data(),
adjustment_reader=self.adj_reader
)
def tearDown(self):
del self.data_portal
@classmethod
def create_algo(cls, code, filename=None, sim_params=None):
def create_algo(self, code, filename=None, sim_params=None):
if sim_params is None:
sim_params = cls.sim_params
sim_params = self.sim_params
return TradingAlgorithm(
script=code,
sim_params=sim_params,
env=cls.env,
env=self.env,
algo_filename=filename
)
@@ -254,10 +182,10 @@ class TestAPIShim(TestCase):
similar) hit the same code paths on the DataPortal.
"""
test_start_minute = self.env.market_minutes_for_day(
self.trading_days[0]
self.sim_params.trading_days[0]
)[1]
test_end_minute = self.env.market_minutes_for_day(
self.trading_days[0]
self.sim_params.trading_days[0]
)[-1]
bar_data = BarData(
self.data_portal,
@@ -450,7 +378,7 @@ class TestAPIShim(TestCase):
warnings.simplefilter("default", ZiplineDeprecationWarning)
sim_params = SimulationParameters(
period_start=self.trading_days[1],
period_start=self.sim_params.trading_days[1],
period_end=self.sim_params.period_end,
capital_base=self.sim_params.capital_base,
data_frequency=self.sim_params.data_frequency,
@@ -495,8 +423,8 @@ class TestAPIShim(TestCase):
warnings.simplefilter("default", ZiplineDeprecationWarning)
sim_params = SimulationParameters(
period_start=self.trading_days[8],
period_end=self.trading_days[-1],
period_start=self.sim_params.trading_days[8],
period_end=self.sim_params.trading_days[-1],
data_frequency="minute",
env=self.env
)
+234 -300
View File
@@ -25,21 +25,27 @@ from unittest import TestCase
import uuid
import warnings
import pandas as pd
from pandas.tseries.tools import normalize_date
from pandas.util.testing import assert_frame_equal
from nose.tools import raises
from nose_parameterized import parameterized
from numpy import full, int32, int64
import pandas as pd
from pandas.util.testing import assert_frame_equal
from six import PY2
import sqlalchemy as sa
from zipline.assets import (
Asset,
Equity,
Future,
AssetDBWriter,
AssetFinder,
AssetFinderCachedEquities,
)
from zipline.assets.synthetic import (
make_commodity_future_info,
make_rotating_equity_info,
make_simple_equity_info,
)
from six import itervalues, integer_types
from toolz import valmap
@@ -53,10 +59,7 @@ from zipline.assets.asset_writer import (
write_version_info,
_futures_defaults,
)
from zipline.assets.asset_db_schema import (
ASSET_DB_VERSION,
_version_table_schema,
)
from zipline.assets.asset_db_schema import ASSET_DB_VERSION
from zipline.assets.asset_db_migrations import (
downgrade
)
@@ -66,20 +69,21 @@ from zipline.errors import (
MultipleSymbolsFound,
RootSymbolNotFound,
AssetDBVersionError,
SidAssignmentError,
SidsNotFound,
SymbolNotFound,
AssetDBImpossibleDowngrade,
)
from zipline.finance.trading import TradingEnvironment, noop_load
from zipline.testing import (
all_subindices,
make_commodity_future_info,
make_rotating_equity_info,
make_simple_equity_info,
empty_assets_db,
tmp_assets_db,
tmp_asset_finder,
)
from zipline.testing.predicates import assert_equal
from zipline.testing.fixtures import (
WithAssetFinder,
ZiplineTestCase,
)
from zipline.utils.tradingcalendar import trading_day
@contextmanager
@@ -280,55 +284,69 @@ class AssetTestCase(TestCase):
'a' < Asset(3)
class TestFuture(TestCase):
class TestFuture(WithAssetFinder, ZiplineTestCase):
@classmethod
def make_futures_info(cls):
return pd.DataFrame.from_dict(
{
2468: {
'symbol': 'OMH15',
'root_symbol': 'OM',
'notice_date': pd.Timestamp('2014-01-20', tz='UTC'),
'expiration_date': pd.Timestamp('2014-02-20', tz='UTC'),
'auto_close_date': pd.Timestamp('2014-01-18', tz='UTC'),
'tick_size': .01,
'multiplier': 500.0,
},
0: {
'symbol': 'CLG06',
'root_symbol': 'CL',
'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
'notice_date': pd.Timestamp('2005-12-20', tz='UTC'),
'expiration_date': pd.Timestamp('2006-01-20', tz='UTC'),
'multiplier': 1.0,
},
},
orient='index',
)
@classmethod
def setUpClass(cls):
cls.future = Future(
2468,
symbol='OMH15',
root_symbol='OM',
notice_date=pd.Timestamp('2014-01-20', tz='UTC'),
expiration_date=pd.Timestamp('2014-02-20', tz='UTC'),
auto_close_date=pd.Timestamp('2014-01-18', tz='UTC'),
tick_size=.01,
multiplier=500
)
cls.future2 = Future(
0,
symbol='CLG06',
root_symbol='CL',
start_date=pd.Timestamp('2005-12-01', tz='UTC'),
notice_date=pd.Timestamp('2005-12-20', tz='UTC'),
expiration_date=pd.Timestamp('2006-01-20', tz='UTC')
)
env = TradingEnvironment(load=noop_load)
env.write_data(futures_identifiers=[TestFuture.future,
TestFuture.future2])
cls.asset_finder = env.asset_finder
def init_class_fixtures(cls):
super(TestFuture, cls).init_class_fixtures()
cls.future = cls.asset_finder.lookup_future_symbol('OMH15')
cls.future2 = cls.asset_finder.lookup_future_symbol('CLG06')
def test_str(self):
strd = self.future.__str__()
strd = str(self.future)
self.assertEqual("Future(2468 [OMH15])", strd)
def test_repr(self):
reprd = self.future.__repr__()
self.assertTrue("Future" in reprd)
self.assertTrue("2468" in reprd)
self.assertTrue("OMH15" in reprd)
self.assertTrue("root_symbol='OM'" in reprd)
self.assertTrue(("notice_date=Timestamp('2014-01-20 00:00:00+0000', "
"tz='UTC')") in reprd)
self.assertTrue("expiration_date=Timestamp('2014-02-20 00:00:00+0000'"
in reprd)
self.assertTrue("auto_close_date=Timestamp('2014-01-18 00:00:00+0000'"
in reprd)
self.assertTrue("tick_size=0.01" in reprd)
self.assertTrue("multiplier=500" in reprd)
reprd = repr(self.future)
self.assertIn("Future", reprd)
self.assertIn("2468", reprd)
self.assertIn("OMH15", reprd)
self.assertIn("root_symbol=%s'OM'" % ('u' if PY2 else ''), reprd)
self.assertIn(
"notice_date=Timestamp('2014-01-20 00:00:00+0000', tz='UTC')",
reprd,
)
self.assertIn(
"expiration_date=Timestamp('2014-02-20 00:00:00+0000'",
reprd,
)
self.assertIn(
"auto_close_date=Timestamp('2014-01-18 00:00:00+0000'",
reprd,
)
self.assertIn("tick_size=0.01", reprd)
self.assertIn("multiplier=500", reprd)
@raises(AssertionError)
def test_reduce(self):
reduced = self.future.__reduce__()
self.assertEqual(Future, reduced[0])
assert_equal(
pickle.loads(pickle.dumps(self.future)).to_dict(),
self.future.to_dict(),
)
def test_to_and_from_dict(self):
dictd = self.future.to_dict()
@@ -378,11 +396,18 @@ class TestFuture(TestCase):
TestFuture.asset_finder.lookup_future_symbol('XXX99')
class AssetFinderTestCase(TestCase):
class AssetFinderTestCase(ZiplineTestCase):
asset_finder_type = AssetFinder
def setUp(self):
self.env = TradingEnvironment(load=noop_load)
self.asset_finder_type = AssetFinder
def write_assets(self, **kwargs):
self._asset_writer.write(**kwargs)
def init_instance_fixtures(self):
super(AssetFinderTestCase, self).init_instance_fixtures()
conn = self.enter_instance_context(empty_assets_db())
self._asset_writer = AssetDBWriter(conn)
self.asset_finder = self.asset_finder_type(conn)
def test_lookup_symbol_delimited(self):
as_of = pd.Timestamp('2013-01-01', tz='UTC')
@@ -399,8 +424,8 @@ class AssetFinderTestCase(TestCase):
for i in range(3)
]
)
self.env.write_data(equities_df=frame)
finder = self.asset_finder_type(self.env.engine)
self.write_assets(equities=frame)
finder = self.asset_finder
asset_0, asset_1, asset_2 = (
finder.retrieve_asset(i) for i in range(3)
)
@@ -423,13 +448,13 @@ class AssetFinderTestCase(TestCase):
)
def test_lookup_symbol_fuzzy(self):
metadata = {
0: {'symbol': 'PRTY_HRD'},
1: {'symbol': 'BRKA'},
2: {'symbol': 'BRK_A'},
}
self.env.write_data(equities_data=metadata)
finder = self.env.asset_finder
metadata = pd.DataFrame.from_records([
{'symbol': 'PRTY_HRD'},
{'symbol': 'BRKA'},
{'symbol': 'BRK_A'},
])
self.write_assets(equities=metadata)
finder = self.asset_finder
dt = pd.Timestamp('2013-01-01', tz='UTC')
# Try combos of looking up PRTYHRD with and without a time or fuzzy
@@ -478,8 +503,8 @@ class AssetFinderTestCase(TestCase):
for i, date in enumerate(dates)
]
)
self.env.write_data(equities_df=df)
finder = self.asset_finder_type(self.env.engine)
self.write_assets(equities=df)
finder = self.asset_finder
for _ in range(2): # Run checks twice to test for caching bugs.
with self.assertRaises(SymbolNotFound):
finder.lookup_symbol('NON_EXISTING', dates[0])
@@ -543,24 +568,22 @@ class AssetFinderTestCase(TestCase):
]
)
self.write_assets(equities=df)
def check(expected_sid, date):
result = finder.lookup_symbol(
result = self.asset_finder.lookup_symbol(
'MULTIPLE', date,
)
self.assertEqual(result.symbol, 'MULTIPLE')
self.assertEqual(result.sid, expected_sid)
with tmp_asset_finder(finder_cls=self.asset_finder_type,
equities=df) as finder:
self.assertIsInstance(finder, self.asset_finder_type)
# Sids 1 and 2 are eligible here. We should get asset 2 because it
# has the later end_date.
check(2, pd.Timestamp('2010-12-31'))
# Sids 1 and 2 are eligible here. We should get asset 2 because it
# has the later end_date.
check(2, pd.Timestamp('2010-12-31'))
# Sids 1, 2, and 3 are eligible here. We should get sid 3 because
# it has a later start_date
check(3, pd.Timestamp('2011-01-01'))
# Sids 1, 2, and 3 are eligible here. We should get sid 3 because
# it has a later start_date
check(3, pd.Timestamp('2011-01-01'))
def test_lookup_generic(self):
"""
@@ -610,8 +633,8 @@ class AssetFinderTestCase(TestCase):
},
]
)
self.env.write_data(equities_df=data)
finder = self.asset_finder_type(self.env.engine)
self.write_assets(equities=data)
finder = self.asset_finder
results, missing = finder.lookup_generic(
['REAL', 1, 'FAKE', 'REAL_BUT_OLD', 'REAL_BUT_IN_THE_FUTURE'],
pd.Timestamp('2013-02-01', tz='UTC'),
@@ -629,97 +652,6 @@ class AssetFinderTestCase(TestCase):
self.assertEqual(missing[0], 'FAKE')
self.assertEqual(missing[1], 'REAL_BUT_IN_THE_FUTURE')
def test_insert_metadata(self):
data = {0: {'start_date': '2014-01-01',
'end_date': '2015-01-01',
'symbol': "PLAY",
'foo_data': "FOO"}}
self.env.write_data(equities_data=data)
finder = self.asset_finder_type(self.env.engine)
# Test proper insertion
equity = finder.retrieve_asset(0)
self.assertIsInstance(equity, Equity)
self.assertEqual('PLAY', equity.symbol)
self.assertEqual(pd.Timestamp('2015-01-01', tz='UTC'),
equity.end_date)
# Test invalid field
with self.assertRaises(AttributeError):
equity.foo_data
def test_consume_metadata(self):
# Test dict consumption
dict_to_consume = {0: {'symbol': 'PLAY'},
1: {'symbol': 'MSFT'}}
self.env.write_data(equities_data=dict_to_consume)
finder = self.asset_finder_type(self.env.engine)
equity = finder.retrieve_asset(0)
self.assertIsInstance(equity, Equity)
self.assertEqual('PLAY', equity.symbol)
# Test dataframe consumption
df = pd.DataFrame(columns=['asset_name', 'exchange'], index=[0, 1])
df['asset_name'][0] = "Dave'N'Busters"
df['exchange'][0] = "NASDAQ"
df['asset_name'][1] = "Microsoft"
df['exchange'][1] = "NYSE"
self.env = TradingEnvironment(load=noop_load)
self.env.write_data(equities_df=df)
finder = self.asset_finder_type(self.env.engine)
self.assertEqual('NASDAQ', finder.retrieve_asset(0).exchange)
self.assertEqual('Microsoft', finder.retrieve_asset(1).asset_name)
def test_consume_asset_as_identifier(self):
# Build some end dates
eq_end = pd.Timestamp('2012-01-01', tz='UTC')
fut_end = pd.Timestamp('2008-01-01', tz='UTC')
# Build some simple Assets
equity_asset = Equity(1, symbol="TESTEQ", end_date=eq_end)
future_asset = Future(200, symbol="TESTFUT", end_date=fut_end)
# Consume the Assets
self.env.write_data(equities_identifiers=[equity_asset],
futures_identifiers=[future_asset])
finder = self.asset_finder_type(self.env.engine)
# Test equality with newly built Assets
self.assertEqual(equity_asset, finder.retrieve_asset(1))
self.assertEqual(future_asset, finder.retrieve_asset(200))
self.assertEqual(eq_end, finder.retrieve_asset(1).end_date)
self.assertEqual(fut_end, finder.retrieve_asset(200).end_date)
def test_sid_assignment(self):
# This metadata does not contain SIDs
metadata = ['PLAY', 'MSFT']
today = normalize_date(pd.Timestamp('2015-07-09', tz='UTC'))
# Write data with sid assignment
self.env.write_data(equities_identifiers=metadata,
allow_sid_assignment=True)
# Verify that Assets were built and different sids were assigned
finder = self.asset_finder_type(self.env.engine)
play = finder.lookup_symbol('PLAY', today)
msft = finder.lookup_symbol('MSFT', today)
self.assertEqual('PLAY', play.symbol)
self.assertIsNotNone(play.sid)
self.assertNotEqual(play.sid, msft.sid)
def test_sid_assignment_failure(self):
# This metadata does not contain SIDs
metadata = ['PLAY', 'MSFT']
# Write data without sid assignment, asserting failure
with self.assertRaises(SidAssignmentError):
self.env.write_data(equities_identifiers=metadata,
allow_sid_assignment=False)
def test_security_dates_warning(self):
# Build an asset with an end_date
@@ -740,16 +672,16 @@ class AssetFinderTestCase(TestCase):
DeprecationWarning))
def test_lookup_future_chain(self):
metadata = {
metadata = pd.DataFrame.from_records([
# Notice day is today, so should be valid.
0: {
{
'symbol': 'ADN15',
'root_symbol': 'AD',
'notice_date': pd.Timestamp('2015-06-14', tz='UTC'),
'expiration_date': pd.Timestamp('2015-08-14', tz='UTC'),
'start_date': pd.Timestamp('2015-01-01', tz='UTC')
},
1: {
{
'symbol': 'ADV15',
'root_symbol': 'AD',
'notice_date': pd.Timestamp('2015-05-14', tz='UTC'),
@@ -757,7 +689,7 @@ class AssetFinderTestCase(TestCase):
'start_date': pd.Timestamp('2015-01-01', tz='UTC')
},
# Starts trading today, so should be valid.
2: {
{
'symbol': 'ADF16',
'root_symbol': 'AD',
'notice_date': pd.Timestamp('2015-11-16', tz='UTC'),
@@ -765,7 +697,7 @@ class AssetFinderTestCase(TestCase):
'start_date': pd.Timestamp('2015-05-14', tz='UTC')
},
# Starts trading in August, so not valid.
3: {
{
'symbol': 'ADX16',
'root_symbol': 'AD',
'notice_date': pd.Timestamp('2015-11-16', tz='UTC'),
@@ -773,7 +705,7 @@ class AssetFinderTestCase(TestCase):
'start_date': pd.Timestamp('2015-08-01', tz='UTC')
},
# Notice date comes after expiration
4: {
{
'symbol': 'ADZ16',
'root_symbol': 'AD',
'notice_date': pd.Timestamp('2016-11-25', tz='UTC'),
@@ -782,15 +714,15 @@ class AssetFinderTestCase(TestCase):
},
# This contract has no start date and also this contract should be
# last in all chains
5: {
{
'symbol': 'ADZ20',
'root_symbol': 'AD',
'notice_date': pd.Timestamp('2020-11-25', tz='UTC'),
'expiration_date': pd.Timestamp('2020-11-16', tz='UTC')
},
}
self.env.write_data(futures_data=metadata)
finder = self.asset_finder_type(self.env.engine)
])
self.write_assets(futures=metadata)
finder = self.asset_finder
dt = pd.Timestamp('2015-05-14', tz='UTC')
dt_2 = pd.Timestamp('2015-10-14', tz='UTC')
dt_3 = pd.Timestamp('2016-11-17', tz='UTC')
@@ -824,7 +756,7 @@ class AssetFinderTestCase(TestCase):
def test_map_identifier_index_to_sids(self):
# Build an empty finder and some Assets
dt = pd.Timestamp('2014-01-01', tz='UTC')
finder = self.asset_finder_type(self.env.engine)
finder = self.asset_finder
asset1 = Equity(1, symbol="AAPL")
asset2 = Equity(2, symbol="GOOG")
asset200 = Future(200, symbol="CLK15")
@@ -844,19 +776,17 @@ class AssetFinderTestCase(TestCase):
def test_compute_lifetimes(self):
num_assets = 4
trading_day = self.env.trading_day
first_start = pd.Timestamp('2015-04-01', tz='UTC')
frame = make_rotating_equity_info(
num_assets=num_assets,
first_start=first_start,
frequency=self.env.trading_day,
frequency=trading_day,
periods_between_starts=3,
asset_lifetime=5
)
self.env.write_data(equities_df=frame)
finder = self.env.asset_finder
self.write_assets(equities=frame)
finder = self.asset_finder
all_dates = pd.date_range(
start=first_start,
@@ -904,12 +834,12 @@ class AssetFinderTestCase(TestCase):
def test_sids(self):
# Ensure that the sids property of the AssetFinder is functioning
self.env.write_data(equities_identifiers=[1, 2, 3])
sids = self.env.asset_finder.sids
self.assertEqual(3, len(sids))
self.assertTrue(1 in sids)
self.assertTrue(2 in sids)
self.assertTrue(3 in sids)
self.write_assets(equities=make_simple_equity_info(
[0, 1, 2],
pd.Timestamp('2014-01-01'),
pd.Timestamp('2014-01-02'),
))
self.assertEqual({0, 1, 2}, set(self.asset_finder.sids))
def test_group_by_type(self):
equities = make_simple_equity_info(
@@ -929,13 +859,17 @@ class AssetFinderTestCase(TestCase):
([0, 2, 3], [7, 10]),
(list(equities.index), list(futures.index)),
]
with tmp_asset_finder(equities=equities, futures=futures) as finder:
for equity_sids, future_sids in queries:
results = finder.group_by_type(equity_sids + future_sids)
self.assertEqual(
results,
{'equity': set(equity_sids), 'future': set(future_sids)},
)
self.write_assets(
equities=equities,
futures=futures,
)
finder = self.asset_finder
for equity_sids, future_sids in queries:
results = finder.group_by_type(equity_sids + future_sids)
self.assertEqual(
results,
{'equity': set(equity_sids), 'future': set(future_sids)},
)
@parameterized.expand([
(Equity, 'retrieve_equities', EquitiesNotFound),
@@ -962,26 +896,30 @@ class AssetFinderTestCase(TestCase):
fail_sids = equity_sids
success_sids = future_sids
with tmp_asset_finder(equities=equities, futures=futures) as finder:
# Run twice to exercise caching.
lookup = getattr(finder, lookup_name)
for _ in range(2):
results = lookup(success_sids)
self.assertIsInstance(results, dict)
self.assertEqual(set(results.keys()), set(success_sids))
self.assertEqual(
valmap(int, results),
dict(zip(success_sids, success_sids)),
)
self.assertEqual(
{type_},
{type(asset) for asset in itervalues(results)},
)
with self.assertRaises(failure_type):
lookup(fail_sids)
with self.assertRaises(failure_type):
# Should fail if **any** of the assets are bad.
lookup([success_sids[0], fail_sids[0]])
self.write_assets(
equities=equities,
futures=futures,
)
finder = self.asset_finder
# Run twice to exercise caching.
lookup = getattr(finder, lookup_name)
for _ in range(2):
results = lookup(success_sids)
self.assertIsInstance(results, dict)
self.assertEqual(set(results.keys()), set(success_sids))
self.assertEqual(
valmap(int, results),
dict(zip(success_sids, success_sids)),
)
self.assertEqual(
{type_},
{type(asset) for asset in itervalues(results)},
)
with self.assertRaises(failure_type):
lookup(fail_sids)
with self.assertRaises(failure_type):
# Should fail if **any** of the assets are bad.
lookup([success_sids[0], fail_sids[0]])
def test_retrieve_all(self):
equities = make_simple_equity_info(
@@ -995,43 +933,46 @@ class AssetFinderTestCase(TestCase):
root_symbols=['CL'],
years=[2014],
)
self.write_assets(
equities=equities,
futures=futures,
)
finder = self.asset_finder
all_sids = finder.sids
self.assertEqual(len(all_sids), len(equities) + len(futures))
queries = [
# Empty Query.
(),
# Only Equities.
tuple(equities.index[:2]),
# Only Futures.
tuple(futures.index[:3]),
# Mixed, all cache misses.
tuple(equities.index[2:]) + tuple(futures.index[3:]),
# Mixed, all cache hits.
tuple(equities.index[2:]) + tuple(futures.index[3:]),
# Everything.
all_sids,
all_sids,
]
for sids in queries:
equity_sids = [i for i in sids if i <= max_equity]
future_sids = [i for i in sids if i > max_equity]
results = finder.retrieve_all(sids)
self.assertEqual(sids, tuple(map(int, results)))
with tmp_asset_finder(equities=equities, futures=futures) as finder:
all_sids = finder.sids
self.assertEqual(len(all_sids), len(equities) + len(futures))
queries = [
# Empty Query.
(),
# Only Equities.
tuple(equities.index[:2]),
# Only Futures.
tuple(futures.index[:3]),
# Mixed, all cache misses.
tuple(equities.index[2:]) + tuple(futures.index[3:]),
# Mixed, all cache hits.
tuple(equities.index[2:]) + tuple(futures.index[3:]),
# Everything.
all_sids,
all_sids,
]
for sids in queries:
equity_sids = [i for i in sids if i <= max_equity]
future_sids = [i for i in sids if i > max_equity]
results = finder.retrieve_all(sids)
self.assertEqual(sids, tuple(map(int, results)))
self.assertEqual(
[Equity for _ in equity_sids] +
[Future for _ in future_sids],
list(map(type, results)),
)
self.assertEqual(
(
list(equities.symbol.loc[equity_sids]) +
list(futures.symbol.loc[future_sids])
),
list(asset.symbol for asset in results),
)
self.assertEqual(
[Equity for _ in equity_sids] +
[Future for _ in future_sids],
list(map(type, results)),
)
self.assertEqual(
(
list(equities.symbol.loc[equity_sids]) +
list(futures.symbol.loc[future_sids])
),
list(asset.symbol for asset in results),
)
@parameterized.expand([
(EquitiesNotFound, 'equity', 'equities'),
@@ -1059,50 +1000,46 @@ class AssetFinderTestCase(TestCase):
class AssetFinderCachedEquitiesTestCase(AssetFinderTestCase):
asset_finder_type = AssetFinderCachedEquities
def setUp(self):
self.env = TradingEnvironment(load=noop_load)
self.asset_finder_type = AssetFinderCachedEquities
def write_assets(self, **kwargs):
super(AssetFinderCachedEquitiesTestCase, self).write_assets(**kwargs)
self.asset_finder.rehash_equities()
class TestFutureChain(TestCase):
class TestFutureChain(WithAssetFinder, ZiplineTestCase):
@classmethod
def setUpClass(cls):
metadata = {
0: {
def make_futures_info(cls):
return pd.DataFrame.from_records([
{
'symbol': 'CLG06',
'root_symbol': 'CL',
'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
'notice_date': pd.Timestamp('2005-12-20', tz='UTC'),
'expiration_date': pd.Timestamp('2006-01-20', tz='UTC')},
1: {
'expiration_date': pd.Timestamp('2006-01-20', tz='UTC'),
},
{
'root_symbol': 'CL',
'symbol': 'CLK06',
'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
'notice_date': pd.Timestamp('2006-03-20', tz='UTC'),
'expiration_date': pd.Timestamp('2006-04-20', tz='UTC')},
2: {
'expiration_date': pd.Timestamp('2006-04-20', tz='UTC'),
},
{
'symbol': 'CLQ06',
'root_symbol': 'CL',
'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
'notice_date': pd.Timestamp('2006-06-20', tz='UTC'),
'expiration_date': pd.Timestamp('2006-07-20', tz='UTC')},
3: {
'expiration_date': pd.Timestamp('2006-07-20', tz='UTC'),
},
{
'symbol': 'CLX06',
'root_symbol': 'CL',
'start_date': pd.Timestamp('2006-02-01', tz='UTC'),
'notice_date': pd.Timestamp('2006-09-20', tz='UTC'),
'expiration_date': pd.Timestamp('2006-10-20', tz='UTC')}
}
env = TradingEnvironment(load=noop_load)
env.write_data(futures_data=metadata)
cls.asset_finder = env.asset_finder
@classmethod
def tearDownClass(cls):
del cls.asset_finder
'expiration_date': pd.Timestamp('2006-10-20', tz='UTC'),
}
])
def test_len(self):
""" Test the __len__ method of FutureChain.
@@ -1310,11 +1247,15 @@ class TestFutureChain(TestCase):
self.assertEqual(codes[key], month_to_cme_code(key))
class TestAssetDBVersioning(TestCase):
class TestAssetDBVersioning(ZiplineTestCase):
def init_instance_fixtures(self):
super(TestAssetDBVersioning, self).init_instance_fixtures()
self.engine = eng = self.enter_instance_context(empty_assets_db())
self.metadata = sa.MetaData(eng, reflect=True)
def test_check_version(self):
env = TradingEnvironment(load=noop_load)
version_table = env.asset_finder.version_info
version_table = self.metadata.tables['version_info']
# This should not raise an error
check_version_info(version_table, ASSET_DB_VERSION)
@@ -1328,9 +1269,7 @@ class TestAssetDBVersioning(TestCase):
check_version_info(version_table, ASSET_DB_VERSION + 1)
def test_write_version(self):
env = TradingEnvironment(load=noop_load)
metadata = sa.MetaData(bind=env.engine)
version_table = _version_table_schema(metadata)
version_table = self.metadata.tables['version_info']
version_table.delete().execute()
# Assert that the version is not present in the table
@@ -1353,17 +1292,14 @@ class TestAssetDBVersioning(TestCase):
write_version_info(version_table, -3)
def test_finder_checks_version(self):
# Create an env and give it a bogus version number
env = TradingEnvironment(load=noop_load)
metadata = sa.MetaData(bind=env.engine)
version_table = _version_table_schema(metadata)
version_table = self.metadata.tables['version_info']
version_table.delete().execute()
write_version_info(version_table, -2)
check_version_info(version_table, -2)
# Assert that trying to build a finder with a bad db raises an error
with self.assertRaises(AssetDBVersionError):
AssetFinder(engine=env.engine)
AssetFinder(engine=self.engine)
# Change the version number of the db to the correct version
version_table.delete().execute()
@@ -1371,17 +1307,16 @@ class TestAssetDBVersioning(TestCase):
check_version_info(version_table, ASSET_DB_VERSION)
# Now that the versions match, this Finder should succeed
AssetFinder(engine=env.engine)
AssetFinder(engine=self.engine)
def test_downgrade(self):
# Attempt to downgrade a current assets db all the way down to v0
env = TradingEnvironment(load=noop_load)
conn = env.engine.connect()
downgrade(env.engine, 0)
conn = self.engine.connect()
downgrade(self.engine, 0)
# Verify that the db version is now 0
metadata = sa.MetaData(conn)
metadata.reflect(bind=env.engine)
metadata.reflect(bind=self.engine)
version_table = metadata.tables['version_info']
check_version_info(version_table, 0)
@@ -1396,6 +1331,5 @@ class TestAssetDBVersioning(TestCase):
def test_impossible_downgrade(self):
# Attempt to downgrade a current assets db to a
# higher-than-current version
env = TradingEnvironment(load=noop_load)
with self.assertRaises(AssetDBImpossibleDowngrade):
downgrade(env.engine, ASSET_DB_VERSION + 5)
downgrade(self.engine, ASSET_DB_VERSION + 5)
+247 -251
View File
@@ -12,24 +12,23 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import TestCase
from testfixtures import TempDirectory
import pandas as pd
import numpy as np
from nose_parameterized import parameterized
import numpy as np
import pandas as pd
from toolz import merge
from zipline._protocol import handle_non_market_minutes
from zipline.data.data_portal import DataPortal
from zipline.data.minute_bars import BcolzMinuteBarWriter, \
US_EQUITIES_MINUTES_PER_DAY, BcolzMinuteBarReader
from zipline.data.us_equity_pricing import BcolzDailyBarReader, \
SQLiteAdjustmentReader, SQLiteAdjustmentWriter
from zipline.finance.trading import TradingEnvironment
from zipline.protocol import BarData
from zipline.testing.core import write_minute_data_for_asset, \
create_daily_df_for_asset, DailyBarWriterFromDataFrames, \
create_mock_adjustments, str_to_seconds, MockDailyBarReader
from zipline.testing import (
MockDailyBarReader,
create_daily_df_for_asset,
create_minute_df_for_asset,
str_to_seconds,
)
from zipline.testing.fixtures import (
WithDataPortal,
ZiplineTestCase,
)
OHLC = ["open", "high", "low", "close"]
OHLCP = OHLC + ["price"]
@@ -44,7 +43,7 @@ field_info = {
}
class TestBarDataBase(TestCase):
class WithBarDataChecks(object):
def assert_same(self, val1, val2):
try:
self.assertEqual(val1, val2)
@@ -89,117 +88,92 @@ class TestBarDataBase(TestCase):
getattr(bar_data, field)
class TestMinuteBarData(TestBarDataBase):
@classmethod
def setUpClass(cls):
cls.tempdir = TempDirectory()
class TestMinuteBarData(WithBarDataChecks,
WithDataPortal,
ZiplineTestCase):
START_DATE = pd.Timestamp('2016-01-05', tz='UTC')
END_DATE = ASSET_FINDER_EQUITY_END_DATE = pd.Timestamp(
'2016-01-07',
tz='UTC',
)
ASSET_FINDER_EQUITY_SIDS = 1, 2, 3, 4, 5
SPLIT_ASSET_SID = 3
ILLIQUID_SPLIT_ASSET_SID = 4
HILARIOUSLY_ILLIQUID_ASSET_SID = 5
@classmethod
def make_minute_bar_data(cls):
# asset1 has trades every minute
# asset2 has trades every 10 minutes
# split_asset trades every minute
# illiquid_split_asset trades every 10 minutes
cls.env = TradingEnvironment()
cls.days = cls.env.days_in_range(
start=pd.Timestamp("2016-01-05", tz='UTC'),
end=pd.Timestamp("2016-01-07", tz='UTC')
return merge(
{
sid: create_minute_df_for_asset(
cls.env,
cls.bcolz_minute_bar_days[0],
cls.bcolz_minute_bar_days[-1],
)
for sid in (1, cls.SPLIT_ASSET_SID)
},
{
sid: create_minute_df_for_asset(
cls.env,
cls.bcolz_minute_bar_days[0],
cls.bcolz_minute_bar_days[-1],
10,
)
for sid in (2, cls.ILLIQUID_SPLIT_ASSET_SID)
},
{
cls.HILARIOUSLY_ILLIQUID_ASSET_SID: create_minute_df_for_asset(
cls.env,
cls.bcolz_minute_bar_days[0],
cls.bcolz_minute_bar_days[-1],
50,
)
},
)
cls.env.write_data(equities_data={
sid: {
'start_date': cls.days[0],
'end_date': cls.days[-1],
'symbol': "ASSET{0}".format(sid)
} for sid in [1, 2, 3, 4, 5]
})
@classmethod
def make_splits_data(cls):
return pd.DataFrame([
{
'effective_date': str_to_seconds("2016-01-06"),
'ratio': 0.5,
'sid': cls.SPLIT_ASSET_SID,
},
{
'effective_date': str_to_seconds("2016-01-06"),
'ratio': 0.5,
'sid': cls.ILLIQUID_SPLIT_ASSET_SID,
},
])
cls.ASSET1 = cls.env.asset_finder.retrieve_asset(1)
cls.ASSET2 = cls.env.asset_finder.retrieve_asset(2)
cls.SPLIT_ASSET = cls.env.asset_finder.retrieve_asset(3)
cls.ILLIQUID_SPLIT_ASSET = cls.env.asset_finder.retrieve_asset(4)
cls.HILARIOUSLY_ILLIQUID_ASSET = cls.env.asset_finder.retrieve_asset(5)
@classmethod
def init_class_fixtures(cls):
super(TestMinuteBarData, cls).init_class_fixtures()
cls.ASSET1 = cls.asset_finder.retrieve_asset(1)
cls.ASSET2 = cls.asset_finder.retrieve_asset(2)
cls.SPLIT_ASSET = cls.asset_finder.retrieve_asset(
cls.SPLIT_ASSET_SID,
)
cls.ILLIQUID_SPLIT_ASSET = cls.asset_finder.retrieve_asset(
cls.ILLIQUID_SPLIT_ASSET_SID,
)
cls.HILARIOUSLY_ILLIQUID_ASSET = cls.asset_finder.retrieve_asset(
cls.HILARIOUSLY_ILLIQUID_ASSET_SID,
)
cls.ASSETS = [cls.ASSET1, cls.ASSET2]
cls.adjustments_reader = cls.create_adjustments_reader()
cls.data_portal = DataPortal(
cls.env,
equity_minute_reader=cls.build_minute_data(),
adjustment_reader=cls.adjustments_reader
)
@classmethod
def tearDownClass(cls):
del cls.data_portal
del cls.adjustments_reader
cls.tempdir.cleanup()
@classmethod
def create_adjustments_reader(cls):
path = create_mock_adjustments(
cls.tempdir,
cls.days,
splits=[{
'effective_date': str_to_seconds("2016-01-06"),
'ratio': 0.5,
'sid': cls.SPLIT_ASSET.sid
}, {
'effective_date': str_to_seconds("2016-01-06"),
'ratio': 0.5,
'sid': cls.ILLIQUID_SPLIT_ASSET.sid
}]
)
return SQLiteAdjustmentReader(path)
@classmethod
def build_minute_data(cls):
market_opens = cls.env.open_and_closes.market_open.loc[cls.days]
market_closes = cls.env.open_and_closes.market_close.loc[cls.days]
writer = BcolzMinuteBarWriter(
cls.days[0],
cls.tempdir.path,
market_opens,
market_closes,
US_EQUITIES_MINUTES_PER_DAY
)
for sid in [cls.ASSET1.sid, cls.SPLIT_ASSET.sid]:
write_minute_data_for_asset(
cls.env,
writer,
cls.days[0],
cls.days[-1],
sid
)
for sid in [cls.ASSET2.sid, cls.ILLIQUID_SPLIT_ASSET.sid]:
write_minute_data_for_asset(
cls.env,
writer,
cls.days[0],
cls.days[-1],
sid,
10
)
write_minute_data_for_asset(
cls.env,
writer,
cls.days[0],
cls.days[-1],
cls.HILARIOUSLY_ILLIQUID_ASSET.sid,
50
)
return BcolzMinuteBarReader(cls.tempdir.path)
def test_minute_before_assets_trading(self):
# grab minutes that include the day before the asset start
minutes = self.env.market_minutes_for_day(
self.env.previous_trading_day(self.days[0])
self.env.previous_trading_day(self.bcolz_minute_bar_days[0])
)
# this entire day is before either asset has started trading
@@ -225,7 +199,9 @@ class TestMinuteBarData(TestBarDataBase):
self.assertTrue(asset_value is pd.NaT)
def test_regular_minute(self):
minutes = self.env.market_minutes_for_day(self.days[0])
minutes = self.env.market_minutes_for_day(
self.bcolz_minute_bar_days[0],
)
for idx, minute in enumerate(minutes):
# day2 has prices
@@ -315,7 +291,9 @@ class TestMinuteBarData(TestBarDataBase):
asset2_value)
def test_minute_of_last_day(self):
minutes = self.env.market_minutes_for_day(self.days[-1])
minutes = self.env.market_minutes_for_day(
self.bcolz_daily_bar_days[-1],
)
# this is the last day the assets exist
for idx, minute in enumerate(minutes):
@@ -326,11 +304,11 @@ class TestMinuteBarData(TestBarDataBase):
def test_minute_after_assets_stopped(self):
minutes = self.env.market_minutes_for_day(
self.env.next_trading_day(self.days[-1])
self.env.next_trading_day(self.bcolz_minute_bar_days[-1])
)
last_trading_minute = \
self.env.market_minutes_for_day(self.days[-1])[-1]
self.env.market_minutes_for_day(self.bcolz_minute_bar_days[-1])[-1]
# this entire day is after both assets have stopped trading
for idx, minute in enumerate(minutes):
@@ -357,7 +335,7 @@ class TestMinuteBarData(TestBarDataBase):
def test_spot_price_is_unadjusted(self):
# verify there is a split for SPLIT_ASSET
splits = self.adjustments_reader.get_adjustments_for_sid(
splits = self.adjustment_reader.get_adjustments_for_sid(
"splits",
self.SPLIT_ASSET.sid
)
@@ -371,7 +349,8 @@ class TestMinuteBarData(TestBarDataBase):
# ... but that's it's not applied when using spot value
minutes = self.env.minutes_for_days_in_range(
start=self.days[0], end=self.days[1]
start=self.bcolz_minute_bar_days[0],
end=self.bcolz_minute_bar_days[1],
)
for idx, minute in enumerate(minutes):
@@ -384,8 +363,12 @@ class TestMinuteBarData(TestBarDataBase):
def test_spot_price_is_adjusted_if_needed(self):
# on cls.days[1], the first 9 minutes of ILLIQUID_SPLIT_ASSET are
# missing. let's get them.
day0_minutes = self.env.market_minutes_for_day(self.days[0])
day1_minutes = self.env.market_minutes_for_day(self.days[1])
day0_minutes = self.env.market_minutes_for_day(
self.bcolz_minute_bar_days[0],
)
day1_minutes = self.env.market_minutes_for_day(
self.bcolz_minute_bar_days[1],
)
for idx, minute in enumerate(day0_minutes[-10:-1]):
bar_data = BarData(self.data_portal, lambda: minute, "minute")
@@ -415,7 +398,7 @@ class TestMinuteBarData(TestBarDataBase):
def test_spot_price_at_midnight(self):
# make sure that if we try to get a minute price at a non-market
# minute, we use the previous market close's timestamp
day = self.days[1]
day = self.bcolz_minute_bar_days[1]
eight_fortyfive_am_eastern = \
pd.Timestamp("{0}-{1}-{2} 8:45".format(
@@ -457,7 +440,9 @@ class TestMinuteBarData(TestBarDataBase):
def test_can_trade_at_midnight(self):
# make sure that if we use `can_trade` at midnight, we don't pretend
# we're in the previous day's last minute
the_day_after = self.env.next_trading_day(self.days[-1])
the_day_after = self.env.next_trading_day(
self.bcolz_minute_bar_days[-1],
)
bar_data = BarData(self.data_portal, lambda: the_day_after, "minute")
@@ -468,7 +453,11 @@ class TestMinuteBarData(TestBarDataBase):
self.assertFalse(bar_data.can_trade(asset))
# but make sure it works when the assets are alive
bar_data2 = BarData(self.data_portal, lambda: self.days[1], "minute")
bar_data2 = BarData(
self.data_portal,
lambda: self.bcolz_minute_bar_days[1],
"minute",
)
for asset in [self.ASSET1, self.HILARIOUSLY_ILLIQUID_ASSET]:
self.assertTrue(bar_data2.can_trade(asset))
@@ -476,14 +465,18 @@ class TestMinuteBarData(TestBarDataBase):
self.assertTrue(bar_data2.can_trade(asset))
def test_is_stale_at_midnight(self):
bar_data = BarData(self.data_portal, lambda: self.days[1], "minute")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_minute_bar_days[1],
"minute",
)
with handle_non_market_minutes(bar_data):
self.assertTrue(bar_data.is_stale(self.HILARIOUSLY_ILLIQUID_ASSET))
def test_overnight_adjustments(self):
# verify there is a split for SPLIT_ASSET
splits = self.adjustments_reader.get_adjustments_for_sid(
splits = self.adjustment_reader.get_adjustments_for_sid(
"splits",
self.SPLIT_ASSET.sid
)
@@ -496,7 +489,7 @@ class TestMinuteBarData(TestBarDataBase):
)
# Current day is 1/06/16
day = self.days[1]
day = self.bcolz_daily_bar_days[1]
eight_fortyfive_am_eastern = \
pd.Timestamp("{0}-{1}-{2} 8:45".format(
day.year, day.month, day.day),
@@ -524,160 +517,135 @@ class TestMinuteBarData(TestBarDataBase):
self.assertEqual(value, expected[field])
class TestDailyBarData(TestBarDataBase):
@classmethod
def setUpClass(cls):
cls.tempdir = TempDirectory()
class TestDailyBarData(WithBarDataChecks,
WithDataPortal,
ZiplineTestCase):
START_DATE = pd.Timestamp('2016-01-05', tz='UTC')
END_DATE = ASSET_FINDER_EQUITY_END_DATE = pd.Timestamp(
'2016-01-08',
tz='UTC',
)
# asset1 has a daily data for each day (1/5, 1/6, 1/7)
# asset2 only has daily data for day2 (1/6)
sids = ASSET_FINDER_EQUITY_SIDS = set(range(1, 9))
cls.env = TradingEnvironment()
cls.days = cls.env.days_in_range(
start=pd.Timestamp("2016-01-05", tz='UTC'),
end=pd.Timestamp("2016-01-08", tz='UTC')
)
cls.env.write_data(equities_data={
sid: {
'start_date': cls.days[0],
'end_date': cls.days[-1],
'symbol': "ASSET{0}".format(sid)
} for sid in [1, 2, 3, 4, 5, 6, 7, 8]
})
cls.ASSET1 = cls.env.asset_finder.retrieve_asset(1)
cls.ASSET2 = cls.env.asset_finder.retrieve_asset(2)
cls.SPLIT_ASSET = cls.env.asset_finder.retrieve_asset(3)
cls.ILLIQUID_SPLIT_ASSET = cls.env.asset_finder.retrieve_asset(4)
cls.MERGER_ASSET = cls.env.asset_finder.retrieve_asset(5)
cls.ILLIQUID_MERGER_ASSET = cls.env.asset_finder.retrieve_asset(6)
cls.DIVIDEND_ASSET = cls.env.asset_finder.retrieve_asset(7)
cls.ILLIQUID_DIVIDEND_ASSET = cls.env.asset_finder.retrieve_asset(8)
cls.ASSETS = [cls.ASSET1, cls.ASSET2]
cls.adjustments_reader = cls.create_adjustments_reader()
cls.data_portal = DataPortal(
cls.env,
equity_daily_reader=cls.build_daily_data(),
adjustment_reader=cls.adjustments_reader
)
SPLIT_ASSET_SID = 3
ILLIQUID_SPLIT_ASSET_SID = 4
MERGER_ASSET_SID = 5
ILLIQUID_MERGER_ASSET_SID = 6
DIVIDEND_ASSET_SID = 7
ILLIQUID_DIVIDEND_ASSET_SID = 8
@classmethod
def tearDownClass(cls):
del cls.data_portal
del cls.adjustments_reader
cls.tempdir.cleanup()
@classmethod
def create_adjustments_reader(cls):
path = cls.tempdir.getpath("test_adjustments.db")
adj_writer = SQLiteAdjustmentWriter(
path,
cls.env.trading_days,
MockDailyBarReader()
)
splits = pd.DataFrame([
def make_splits_data(cls):
return pd.DataFrame.from_records([
{
'effective_date': str_to_seconds("2016-01-06"),
'ratio': 0.5,
'sid': cls.SPLIT_ASSET.sid
'sid': cls.SPLIT_ASSET_SID,
},
{
'effective_date': str_to_seconds("2016-01-07"),
'ratio': 0.5,
'sid': cls.ILLIQUID_SPLIT_ASSET.sid
}
'sid': cls.ILLIQUID_SPLIT_ASSET_SID,
},
])
mergers = pd.DataFrame([
@classmethod
def make_mergers_data(cls):
return pd.DataFrame.from_records([
{
'effective_date': str_to_seconds("2016-01-06"),
'effective_date': str_to_seconds('2016-01-06'),
'ratio': 0.5,
'sid': cls.MERGER_ASSET.sid
'sid': cls.MERGER_ASSET_SID,
},
{
'effective_date': str_to_seconds("2016-01-07"),
'effective_date': str_to_seconds('2016-01-07'),
'ratio': 0.6,
'sid': cls.ILLIQUID_MERGER_ASSET.sid
'sid': cls.ILLIQUID_MERGER_ASSET_SID,
}
])
# we're using a fake daily reader in the adjustments writer which
# returns every daily price as 100, so dividend amounts of 2.0 and 4.0
# correspond to 2% and 4% dividends, respectively.
dividends = pd.DataFrame([
@classmethod
def make_dividends_data(cls):
return pd.DataFrame.from_records([
{
# only care about ex date, the other dates don't matter here
'ex_date':
pd.Timestamp("2016-01-06", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-06', tz='UTC').to_datetime64(),
'record_date':
pd.Timestamp("2016-01-06", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-06', tz='UTC').to_datetime64(),
'declared_date':
pd.Timestamp("2016-01-06", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-06', tz='UTC').to_datetime64(),
'pay_date':
pd.Timestamp("2016-01-06", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-06', tz='UTC').to_datetime64(),
'amount': 2.0,
'sid': cls.DIVIDEND_ASSET.sid
'sid': cls.DIVIDEND_ASSET_SID,
},
{
'ex_date':
pd.Timestamp("2016-01-07", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-07', tz='UTC').to_datetime64(),
'record_date':
pd.Timestamp("2016-01-07", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-07', tz='UTC').to_datetime64(),
'declared_date':
pd.Timestamp("2016-01-07", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-07', tz='UTC').to_datetime64(),
'pay_date':
pd.Timestamp("2016-01-07", tz='UTC').to_datetime64(),
pd.Timestamp('2016-01-07', tz='UTC').to_datetime64(),
'amount': 4.0,
'sid': cls.ILLIQUID_DIVIDEND_ASSET.sid
'sid': cls.ILLIQUID_DIVIDEND_ASSET_SID,
}],
columns=['ex_date',
'record_date',
'declared_date',
'pay_date',
'amount',
'sid']
columns=[
'ex_date',
'record_date',
'declared_date',
'pay_date',
'amount',
'sid',
]
)
adj_writer.write(splits, mergers, dividends)
return SQLiteAdjustmentReader(path)
@classmethod
def make_adjustment_writer_daily_bar_reader(cls):
return MockDailyBarReader()
@classmethod
def build_daily_data(cls):
path = cls.tempdir.getpath("testdaily.bcolz")
def make_daily_bar_data(cls):
for sid in cls.sids:
yield sid, create_daily_df_for_asset(
cls.env,
cls.bcolz_daily_bar_days[0],
cls.bcolz_daily_bar_days[-1],
interval=2 - sid % 2
)
dfs = {
1: create_daily_df_for_asset(cls.env, cls.days[0], cls.days[-1]),
2: create_daily_df_for_asset(
cls.env, cls.days[0], cls.days[-1], interval=2
),
3: create_daily_df_for_asset(cls.env, cls.days[0], cls.days[-1]),
4: create_daily_df_for_asset(
cls.env, cls.days[0], cls.days[-1], interval=2
),
5: create_daily_df_for_asset(cls.env, cls.days[0], cls.days[-1]),
6: create_daily_df_for_asset(
cls.env, cls.days[0], cls.days[-1], interval=2
),
7: create_daily_df_for_asset(cls.env, cls.days[0], cls.days[-1]),
8: create_daily_df_for_asset(
cls.env, cls.days[0], cls.days[-1], interval=2
),
}
@classmethod
def init_class_fixtures(cls):
super(TestDailyBarData, cls).init_class_fixtures()
daily_writer = DailyBarWriterFromDataFrames(dfs)
daily_writer.write(path, cls.days, dfs)
return BcolzDailyBarReader(path)
cls.ASSET1 = cls.asset_finder.retrieve_asset(1)
cls.ASSET2 = cls.asset_finder.retrieve_asset(2)
cls.SPLIT_ASSET = cls.asset_finder.retrieve_asset(
cls.SPLIT_ASSET_SID,
)
cls.ILLIQUID_SPLIT_ASSET = cls.asset_finder.retrieve_asset(
cls.ILLIQUID_SPLIT_ASSET_SID,
)
cls.MERGER_ASSET = cls.asset_finder.retrieve_asset(
cls.MERGER_ASSET_SID,
)
cls.ILLIQUID_MERGER_ASSET = cls.asset_finder.retrieve_asset(
cls.ILLIQUID_MERGER_ASSET_SID,
)
cls.DIVIDEND_ASSET = cls.asset_finder.retrieve_asset(
cls.DIVIDEND_ASSET_SID,
)
cls.ILLIQUID_DIVIDEND_ASSET = cls.asset_finder.retrieve_asset(
cls.ILLIQUID_DIVIDEND_ASSET_SID,
)
cls.ASSETS = [cls.ASSET1, cls.ASSET2]
def test_day_before_assets_trading(self):
# use the day before self.days[0]
day = self.env.previous_trading_day(self.days[0])
# use the day before self.bcolz_daily_bar_days[0]
day = self.env.previous_trading_day(self.bcolz_daily_bar_days[0])
bar_data = BarData(self.data_portal, lambda: day, "daily")
self.check_internal_consistency(bar_data)
@@ -700,8 +668,12 @@ class TestDailyBarData(TestBarDataBase):
self.assertTrue(asset_value is pd.NaT)
def test_semi_active_day(self):
# on self.days[0], only asset1 has data
bar_data = BarData(self.data_portal, lambda: self.days[0], "daily")
# on self.bcolz_daily_bar_days[0], only asset1 has data
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[0],
"daily",
)
self.check_internal_consistency(bar_data)
self.assertTrue(bar_data.can_trade(self.ASSET1))
@@ -719,7 +691,7 @@ class TestDailyBarData(TestBarDataBase):
self.assertEqual(2, bar_data.current(self.ASSET1, "close"))
self.assertEqual(200, bar_data.current(self.ASSET1, "volume"))
self.assertEqual(2, bar_data.current(self.ASSET1, "price"))
self.assertEqual(self.days[0],
self.assertEqual(self.bcolz_daily_bar_days[0],
bar_data.current(self.ASSET1, "last_traded"))
for field in OHLCP:
@@ -732,10 +704,14 @@ class TestDailyBarData(TestBarDataBase):
)
def test_fully_active_day(self):
bar_data = BarData(self.data_portal, lambda: self.days[1], "daily")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[1],
"daily",
)
self.check_internal_consistency(bar_data)
# on self.days[1], both assets have data
# on self.bcolz_daily_bar_days[1], both assets have data
for asset in self.ASSETS:
self.assertTrue(bar_data.can_trade(asset))
self.assertFalse(bar_data.is_stale(asset))
@@ -747,12 +723,16 @@ class TestDailyBarData(TestBarDataBase):
self.assertEqual(300, bar_data.current(asset, "volume"))
self.assertEqual(3, bar_data.current(asset, "price"))
self.assertEqual(
self.days[1],
self.bcolz_daily_bar_days[1],
bar_data.current(asset, "last_traded")
)
def test_last_active_day(self):
bar_data = BarData(self.data_portal, lambda: self.days[-1], "daily")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[-1],
"daily",
)
self.check_internal_consistency(bar_data)
for asset in self.ASSETS:
@@ -768,7 +748,7 @@ class TestDailyBarData(TestBarDataBase):
def test_after_assets_dead(self):
# both assets end on self.day[-1], so let's try the next day
next_day = self.env.next_trading_day(self.days[-1])
next_day = self.env.next_trading_day(self.bcolz_daily_bar_days[-1])
bar_data = BarData(self.data_portal, lambda: next_day, "daily")
self.check_internal_consistency(bar_data)
@@ -785,9 +765,9 @@ class TestDailyBarData(TestBarDataBase):
last_traded_dt = bar_data.current(asset, "last_traded")
if asset == self.ASSET1:
self.assertEqual(self.days[-2], last_traded_dt)
self.assertEqual(self.bcolz_daily_bar_days[-2], last_traded_dt)
else:
self.assertEqual(self.days[1], last_traded_dt)
self.assertEqual(self.bcolz_daily_bar_days[1], last_traded_dt)
@parameterized.expand([
("split", 2, 3, 3, 1.5),
@@ -808,7 +788,7 @@ class TestDailyBarData(TestBarDataBase):
("ILLIQUID_" + adjustment_type.upper() + "_ASSET")
)
# verify there is an adjustment for liquid_asset
adjustments = self.adjustments_reader.get_adjustments_for_sid(
adjustments = self.adjustment_reader.get_adjustments_for_sid(
table_name,
liquid_asset.sid
)
@@ -821,12 +801,20 @@ class TestDailyBarData(TestBarDataBase):
)
# ... but that's it's not applied when using spot value
bar_data = BarData(self.data_portal, lambda: self.days[0], "daily")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[0],
"daily",
)
self.assertEqual(
liquid_day_0_price,
bar_data.current(liquid_asset, "price")
)
bar_data = BarData(self.data_portal, lambda: self.days[1], "daily")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[1],
"daily",
)
self.assertEqual(
liquid_day_1_price,
bar_data.current(liquid_asset, "price")
@@ -834,12 +822,20 @@ class TestDailyBarData(TestBarDataBase):
# ... except when we have to forward fill across a day boundary
# ILLIQUID_ASSET has no data on days 0 and 2, and a split on day 2
bar_data = BarData(self.data_portal, lambda: self.days[1], "daily")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[1],
"daily",
)
self.assertEqual(
illiquid_day_0_price, bar_data.current(illiquid_asset, "price")
)
bar_data = BarData(self.data_portal, lambda: self.days[2], "daily")
bar_data = BarData(
self.data_portal,
lambda: self.bcolz_daily_bar_days[2],
"daily",
)
# 3 (price from previous day) * 0.5 (split ratio)
self.assertAlmostEqual(
+77 -93
View File
@@ -12,73 +12,66 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from unittest import TestCase
from datetime import timedelta
import numpy as np
import pandas as pd
from testfixtures import TempDirectory
from zipline.data.us_equity_pricing import SQLiteAdjustmentWriter, \
SQLiteAdjustmentReader
from zipline.data.data_portal import DataPortal
from zipline.errors import (
BenchmarkAssetNotAvailableTooEarly,
BenchmarkAssetNotAvailableTooLate,
InvalidBenchmarkAsset)
from zipline.finance.trading import TradingEnvironment
from zipline.sources.benchmark_source import BenchmarkSource
from zipline.utils import factory
from zipline.testing.core import create_data_portal, write_minute_data, \
create_empty_splits_mergers_frame
from .test_perf_tracking import MockDailyBarSpotReader
from zipline.testing import (
MockDailyBarReader,
create_minute_bar_data,
tmp_bcolz_minute_bar_reader,
)
from zipline.testing.fixtures import (
WithDataPortal,
WithSimParams,
ZiplineTestCase,
)
class TestBenchmark(TestCase):
class TestBenchmark(WithDataPortal, WithSimParams, ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
END_DATE = pd.Timestamp('2006-12-29', tz='utc')
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.tempdir = TempDirectory()
cls.sim_params = factory.create_simulation_parameters()
cls.env.write_data(equities_data={
1: {
"start_date": cls.sim_params.trading_days[0],
"end_date": cls.sim_params.trading_days[-1] + timedelta(days=1)
def make_equity_info(cls):
return pd.DataFrame.from_dict(
{
1: {
"start_date": cls.START_DATE,
"end_date": cls.END_DATE + pd.Timedelta(days=1)
},
2: {
"start_date": cls.START_DATE,
"end_date": cls.END_DATE + pd.Timedelta(days=1)
},
3: {
"start_date": pd.Timestamp('2006-05-26', tz='utc'),
"end_date": pd.Timestamp('2006-08-09', tz='utc')
},
4: {
"start_date": cls.START_DATE,
"end_date": cls.END_DATE + pd.Timedelta(days=1)
},
},
2: {
"start_date": cls.sim_params.trading_days[0],
"end_date": cls.sim_params.trading_days[-1] + timedelta(days=1)
},
3: {
"start_date": cls.sim_params.trading_days[100],
"end_date": cls.sim_params.trading_days[-100]
},
4: {
"start_date": cls.sim_params.trading_days[0],
"end_date": cls.sim_params.trading_days[-1] + timedelta(days=1)
}
orient='index',
)
})
@classmethod
def make_adjustment_writer_daily_bar_reader(cls):
return MockDailyBarReader()
dbpath = os.path.join(cls.tempdir.path, "adjustments.db")
writer = SQLiteAdjustmentWriter(dbpath, cls.env.trading_days,
MockDailyBarSpotReader())
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([], dtype=np.uint32),
'amount': np.array([], dtype=np.float64),
'declared_date': np.array([], dtype='datetime64[ns]'),
'ex_date': np.array([], dtype='datetime64[ns]'),
'pay_date': np.array([], dtype='datetime64[ns]'),
'record_date': np.array([], dtype='datetime64[ns]'),
})
@classmethod
def make_stock_dividends_data(cls):
declared_date = cls.sim_params.trading_days[45]
ex_date = cls.sim_params.trading_days[50]
record_date = pay_date = cls.sim_params.trading_days[55]
stock_dividends = pd.DataFrame({
return pd.DataFrame({
'sid': np.array([4], dtype=np.uint32),
'payment_sid': np.array([5], dtype=np.uint32),
'ratio': np.array([2], dtype=np.float64),
@@ -87,22 +80,6 @@ class TestBenchmark(TestCase):
'record_date': np.array([record_date], dtype='datetime64[ns]'),
'pay_date': np.array([pay_date], dtype='datetime64[ns]'),
})
writer.write(splits, mergers, dividends,
stock_dividends=stock_dividends)
cls.data_portal = create_data_portal(
cls.env,
cls.tempdir,
cls.sim_params,
[1, 2, 3, 4],
adjustment_reader=SQLiteAdjustmentReader(dbpath)
)
@classmethod
def tearDownClass(cls):
del cls.data_portal
del cls.env
cls.tempdir.cleanup()
def test_normal(self):
days_to_use = self.sim_params.trading_days[1:]
@@ -162,37 +139,44 @@ class TestBenchmark(TestCase):
self.sim_params.trading_days[5]
)
path = write_minute_data(
tmp_reader = tmp_bcolz_minute_bar_reader(
self.env,
self.tempdir,
minutes,
[2]
self.env.trading_days,
create_minute_bar_data(minutes, [2]),
)
self.data_portal._minutes_equities_path = path
source = BenchmarkSource(
2,
self.env,
self.sim_params.trading_days,
self.data_portal
)
days_to_use = self.sim_params.trading_days
# first value should be 0.0, coming from daily data
self.assertAlmostEquals(0.0, source.get_value(days_to_use[0]))
manually_calculated = self.data_portal.get_history_window(
[2], days_to_use[-1], len(days_to_use), "1d", "close"
)[2].pct_change()
for idx, day in enumerate(days_to_use[1:]):
self.assertEqual(
source.get_value(day),
manually_calculated[idx + 1]
with tmp_reader as reader:
data_portal = DataPortal(
self.env,
equity_minute_reader=reader,
equity_daily_reader=self.bcolz_daily_bar_reader,
adjustment_reader=self.adjustment_reader,
)
source = BenchmarkSource(
2,
self.env,
self.sim_params.trading_days,
data_portal
)
days_to_use = self.sim_params.trading_days
# first value should be 0.0, coming from daily data
self.assertAlmostEquals(0.0, source.get_value(days_to_use[0]))
manually_calculated = data_portal.get_history_window(
[2], days_to_use[-1],
len(days_to_use),
"1d",
"close",
)[2].pct_change()
for idx, day in enumerate(days_to_use[1:]):
self.assertEqual(
source.get_value(day),
manually_calculated[idx + 1]
)
def test_no_stock_dividends_allowed(self):
# try to use sid(4) as benchmark, should blow up due to the presence
# of a stock dividend
+36 -79
View File
@@ -12,16 +12,10 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from nose_parameterized import parameterized
from unittest import TestCase
from testfixtures import TempDirectory
import pandas as pd
import zipline.utils.factory as factory
from zipline.finance import trading
from zipline.finance.blotter import Blotter
from zipline.finance.order import ORDER_STATUS
from zipline.finance.execution import (
@@ -31,89 +25,52 @@ from zipline.finance.execution import (
StopOrder,
)
from zipline.testing import(
setup_logger,
teardown_logger,
)
from zipline.gens.sim_engine import DAY_END, BAR
from zipline.finance.cancel_policy import EODCancel, NeverCancel
from zipline.finance.slippage import DEFAULT_VOLUME_SLIPPAGE_BAR_LIMIT, \
FixedSlippage
from .utils.daily_bar_writer import DailyBarWriterFromDataFrames
from zipline.data.us_equity_pricing import BcolzDailyBarReader
from zipline.data.data_portal import DataPortal
from zipline.finance.slippage import (
DEFAULT_VOLUME_SLIPPAGE_BAR_LIMIT,
FixedSlippage,
)
from zipline.protocol import BarData
from zipline.testing.fixtures import (
WithDataPortal,
WithLogger,
WithSimParams,
ZiplineTestCase,
)
class BlotterTestCase(TestCase):
class BlotterTestCase(WithLogger,
WithDataPortal,
WithSimParams,
ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-05', tz='utc')
END_DATE = pd.Timestamp('2006-01-06', tz='utc')
ASSET_FINDER_EQUITY_SIDS = 24, 25
@classmethod
def setUpClass(cls):
setup_logger(cls)
cls.env = trading.TradingEnvironment()
cls.sim_params = factory.create_simulation_parameters(
start=pd.Timestamp("2006-01-05", tz='UTC'),
end=pd.Timestamp("2006-01-06", tz='UTC')
)
cls.env.write_data(equities_data={
24: {
'start_date': cls.sim_params.trading_days[0],
'end_date': cls.env.next_trading_day(
cls.sim_params.trading_days[-1]
)
def make_daily_bar_data(cls):
yield 24, pd.DataFrame(
{
'open': [50, 50],
'high': [50, 50],
'low': [50, 50],
'close': [50, 50],
'volume': [100, 400],
},
25: {
'start_date': cls.sim_params.trading_days[0],
'end_date': cls.env.next_trading_day(
cls.sim_params.trading_days[-1]
)
}
})
cls.tempdir = TempDirectory()
assets = {
24: pd.DataFrame({
"open": [50, 50],
"high": [50, 50],
"low": [50, 50],
"close": [50, 50],
"volume": [100, 400],
"day": [day.value for day in cls.sim_params.trading_days]
}),
25: pd.DataFrame({
"open": [50, 50],
"high": [50, 50],
"low": [50, 50],
"close": [50, 50],
"volume": [100, 400],
"day": [day.value for day in cls.sim_params.trading_days]
})
}
path = os.path.join(cls.tempdir.path, "tempdata.bcolz")
DailyBarWriterFromDataFrames(assets).write(
path,
cls.sim_params.trading_days,
assets
index=cls.sim_params.trading_days,
)
equity_daily_reader = BcolzDailyBarReader(path)
cls.data_portal = DataPortal(
cls.env,
equity_daily_reader=equity_daily_reader,
yield 25, pd.DataFrame(
{
'open': [50, 50],
'high': [50, 50],
'low': [50, 50],
'close': [50, 50],
'volume': [100, 400],
},
index=cls.sim_params.trading_days,
)
@classmethod
def tearDownClass(cls):
del cls.env
cls.tempdir.cleanup()
teardown_logger(cls)
@parameterized.expand([(MarketOrder(), None, None),
(LimitOrder(10), 10, None),
(StopOrder(10), None, 10),
+9 -37
View File
@@ -12,56 +12,28 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pandas as pd
from unittest import TestCase
from testfixtures import TempDirectory
from zipline.finance.trading import TradingEnvironment
from zipline.test_algorithms import (
ExceptionAlgorithm,
DivByZeroAlgorithm,
SetPortfolioAlgorithm,
)
from zipline.testing import (
setup_logger,
teardown_logger
from zipline.testing.fixtures import (
WithDataPortal,
WithSimParams,
ZiplineTestCase,
)
import zipline.utils.factory as factory
from zipline.testing.core import create_data_portal
DEFAULT_TIMEOUT = 15 # seconds
EXTENDED_TIMEOUT = 90
class ExceptionTestCase(TestCase):
class ExceptionTestCase(WithDataPortal, WithSimParams, ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
START_DATE = pd.Timestamp('2006-01-07', tz='utc')
@classmethod
def setUpClass(cls):
cls.sid = 133
cls.env = TradingEnvironment()
cls.env.write_data(equities_identifiers=[cls.sid])
cls.tempdir = TempDirectory()
cls.sim_params = factory.create_simulation_parameters(
num_days=4,
env=cls.env
)
cls.data_portal = create_data_portal(
env=cls.env,
tempdir=cls.tempdir,
sim_params=cls.sim_params,
sids=[cls.sid]
)
setup_logger(cls)
@classmethod
def tearDownClass(cls):
del cls.env
cls.tempdir.cleanup()
teardown_logger(cls)
sid, = ASSET_FINDER_EQUITY_SIDS = 133,
def test_exception_in_handle_data(self):
algo = ExceptionAlgorithm('handle_data',
+4 -15
View File
@@ -12,30 +12,25 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import TestCase
from nose_parameterized import parameterized
from six.moves import range
from zipline.errors import(
BadOrderParameters
)
from zipline.finance.execution import (
LimitOrder,
MarketOrder,
StopLimitOrder,
StopOrder,
)
from zipline.testing import(
setup_logger,
teardown_logger,
from zipline.testing.fixtures import (
WithLogger,
ZiplineTestCase,
)
class ExecutionStyleTestCase(TestCase):
class ExecutionStyleTestCase(WithLogger, ZiplineTestCase):
"""
Tests for zipline ExecutionStyle classes.
"""
@@ -75,12 +70,6 @@ class ExecutionStyleTestCase(TestCase):
(ArbitraryObject(),),
]
def setUp(self):
setup_logger(self)
def tearDown(self):
teardown_logger(self)
@parameterized.expand(INVALID_PRICES)
def test_invalid_prices(self, price):
"""
+111 -85
View File
@@ -12,120 +12,83 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import TestCase
from nose_parameterized import parameterized
import pandas as pd
import numpy as np
import responses
from mock import patch
from zipline import TradingAlgorithm
from zipline.errors import UnsupportedOrderParameters
from zipline.finance.trading import TradingEnvironment
from zipline.sources.requests_csv import mask_requests_args
from zipline.utils import factory
from zipline.testing.core import FetcherDataPortal
from zipline.testing import FetcherDataPortal
from zipline.testing.fixtures import (
WithResponses,
WithSimParams,
ZiplineTestCase,
)
from .resources.fetcher_inputs.fetcher_test_data import (
MULTI_SIGNAL_CSV_DATA,
AAPL_CSV_DATA,
AAPL_MINUTE_CSV_DATA,
IBM_CSV_DATA,
ANNUAL_AAPL_CSV_DATA,
AAPL_IBM_CSV_DATA,
AAPL_MINUTE_CSV_DATA,
CPIAUCSL_DATA,
PALLADIUM_DATA,
FETCHER_ALTERNATE_COLUMN_HEADER,
FETCHER_UNIVERSE_DATA,
FETCHER_UNIVERSE_DATA_TICKER_COLUMN,
MULTI_SIGNAL_CSV_DATA,
NON_ASSET_FETCHER_UNIVERSE_DATA,
FETCHER_UNIVERSE_DATA_TICKER_COLUMN, FETCHER_ALTERNATE_COLUMN_HEADER)
PALLADIUM_DATA,
)
class FetcherTestCase(TestCase):
class FetcherTestCase(WithResponses,
WithSimParams,
ZiplineTestCase):
@classmethod
def setUpClass(cls):
responses.start()
responses.add(responses.GET,
'https://fake.urls.com/aapl_minute_csv_data.csv',
body=AAPL_MINUTE_CSV_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/aapl_csv_data.csv',
body=AAPL_CSV_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/multi_signal_csv_data.csv',
body=MULTI_SIGNAL_CSV_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/cpiaucsl_data.csv',
body=CPIAUCSL_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/ibm_csv_data.csv',
body=IBM_CSV_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/aapl_ibm_csv_data.csv',
body=AAPL_IBM_CSV_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/palladium_data.csv',
body=PALLADIUM_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/fetcher_universe_data.csv',
body=FETCHER_UNIVERSE_DATA, content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/bad_fetcher_universe_data.csv',
body=NON_ASSET_FETCHER_UNIVERSE_DATA,
content_type='text/csv')
responses.add(responses.GET,
'https://fake.urls.com/annual_aapl_csv_data.csv',
body=ANNUAL_AAPL_CSV_DATA, content_type='text/csv')
cls.sim_params = factory.create_simulation_parameters()
cls.env = TradingEnvironment()
cls.env.write_data(
equities_data={
def make_equity_info(cls):
return pd.DataFrame.from_dict(
{
24: {
"start_date": pd.Timestamp("2006-01-01", tz='UTC'),
"end_date": pd.Timestamp("2007-01-01", tz='UTC'),
'symbol': "AAPL",
"asset_type": "equity",
"exchange": "nasdaq"
'start_date': pd.Timestamp('2006-01-01', tz='UTC'),
'end_date': pd.Timestamp('2007-01-01', tz='UTC'),
'symbol': 'AAPL',
'asset_type': 'equity',
'exchange': 'nasdaq'
},
3766: {
"start_date": pd.Timestamp("2006-01-01", tz='UTC'),
"end_date": pd.Timestamp("2007-01-01", tz='UTC'),
'symbol': "IBM",
"asset_type": "equity",
"exchange": "nasdaq"
'start_date': pd.Timestamp('2006-01-01', tz='UTC'),
'end_date': pd.Timestamp('2007-01-01', tz='UTC'),
'symbol': 'IBM',
'asset_type': 'equity',
'exchange': 'nasdaq'
},
5061: {
"start_date": pd.Timestamp("2006-01-01", tz='UTC'),
"end_date": pd.Timestamp("2007-01-01", tz='UTC'),
'symbol': "MSFT",
"asset_type": "equity",
"exchange": "nasdaq"
'start_date': pd.Timestamp('2006-01-01', tz='UTC'),
'end_date': pd.Timestamp('2007-01-01', tz='UTC'),
'symbol': 'MSFT',
'asset_type': 'equity',
'exchange': 'nasdaq'
},
14848: {
"start_date": pd.Timestamp("2006-01-01", tz='UTC'),
"end_date": pd.Timestamp("2007-01-01", tz='UTC'),
'symbol': "YHOO",
"asset_type": "equity",
"exchange": "nasdaq"
'start_date': pd.Timestamp('2006-01-01', tz='UTC'),
'end_date': pd.Timestamp('2007-01-01', tz='UTC'),
'symbol': 'YHOO',
'asset_type': 'equity',
'exchange': 'nasdaq'
},
25317: {
"start_date": pd.Timestamp("2006-01-01", tz='UTC'),
"end_date": pd.Timestamp("2007-01-01", tz='UTC'),
'symbol': "DELL",
"asset_type": "equity",
"exchange": "nasdaq"
'start_date': pd.Timestamp('2006-01-01', tz='UTC'),
'end_date': pd.Timestamp('2007-01-01', tz='UTC'),
'symbol': 'DELL',
'asset_type': 'equity',
'exchange': 'nasdaq'
}
}
},
orient='index',
)
@classmethod
def tearDownClass(cls):
responses.stop()
responses.reset()
def run_algo(self, code, sim_params=None, data_frequency="daily"):
if sim_params is None:
sim_params = self.sim_params
@@ -141,7 +104,14 @@ class FetcherTestCase(TestCase):
return results
def test_fetch_minute_granularity(self):
def test_minutely_fetcher(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/aapl_minute_csv_data.csv',
body=AAPL_MINUTE_CSV_DATA,
content_type='text/csv',
)
sim_params = factory.create_simulation_parameters(
start=pd.Timestamp("2006-01-03", tz='UTC'),
end=pd.Timestamp("2006-01-10", tz='UTC'),
@@ -195,6 +165,13 @@ def handle_data(context, data):
np.testing.assert_array_equal([4] * 780, signal[1560:])
def test_fetch_csv_with_multi_symbols(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/multi_signal_csv_data.csv',
body=MULTI_SIGNAL_CSV_DATA,
content_type='text/csv',
)
results = self.run_algo(
"""
from zipline.api import fetch_csv, record, sid
@@ -212,6 +189,13 @@ def handle_data(context, data):
self.assertEqual(5, results["dell_signal"].iloc[-1])
def test_fetch_csv_with_pure_signal_file(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/cpiaucsl_data.csv',
body=CPIAUCSL_DATA,
content_type='text/csv',
)
results = self.run_algo(
"""
from zipline.api import fetch_csv, sid, record
@@ -237,6 +221,13 @@ def handle_data(context, data):
self.assertEqual(results["cpi"][-1], 203.1)
def test_algo_fetch_csv(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/aapl_csv_data.csv',
body=AAPL_CSV_DATA,
content_type='text/csv',
)
results = self.run_algo(
"""
from zipline.api import fetch_csv, record, sid
@@ -262,6 +253,13 @@ def handle_data(context, data):
self.assertEqual(24, results["price"][-1]) # fake value
def test_algo_fetch_csv_with_extra_symbols(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/aapl_ibm_csv_data.csv',
body=AAPL_IBM_CSV_DATA,
content_type='text/csv',
)
results = self.run_algo(
"""
from zipline.api import fetch_csv, record, sid
@@ -293,6 +291,13 @@ def handle_data(context, data):
("without date", "usecols=['Value']"),
("with date", "usecols=('Value', 'Date')")])
def test_usecols(self, testname, usecols):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/cpiaucsl_data.csv',
body=CPIAUCSL_DATA,
content_type='text/csv',
)
code = """
from zipline.api import fetch_csv, sid, record
@@ -442,6 +447,13 @@ def handle_data(context, data):
self.assertEqual(4, results["sid_count"].iloc[2])
def test_fetcher_universe_non_security_return(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/bad_fetcher_universe_data.csv',
body=NON_ASSET_FETCHER_UNIVERSE_DATA,
content_type='text/csv',
)
sim_params = factory.create_simulation_parameters(
start=pd.Timestamp("2006-01-09", tz='UTC'),
end=pd.Timestamp("2006-01-10", tz='UTC')
@@ -465,6 +477,13 @@ def handle_data(context, data):
)
def test_order_against_data(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/palladium_data.csv',
body=PALLADIUM_DATA,
content_type='text/csv',
)
with self.assertRaises(UnsupportedOrderParameters):
self.run_algo("""
from zipline.api import fetch_csv, order, sid
@@ -486,6 +505,13 @@ def handle_data(context, data):
""")
def test_fetcher_universe_minute(self):
self.responses.add(
self.responses.GET,
'https://fake.urls.com/fetcher_universe_data.csv',
body=FETCHER_UNIVERSE_DATA,
content_type='text/csv',
)
sim_params = factory.create_simulation_parameters(
start=pd.Timestamp("2006-01-09", tz='UTC'),
end=pd.Timestamp("2006-01-11", tz='UTC'),
+46 -75
View File
@@ -18,8 +18,6 @@ Tests for the zipline.finance package
"""
from datetime import datetime, timedelta
import os
from unittest import TestCase
from nose.tools import timed
import numpy as np
@@ -28,23 +26,27 @@ import pytz
from six.moves import range
from testfixtures import TempDirectory
from zipline.assets.synthetic import make_simple_equity_info
from zipline.finance.blotter import Blotter
from zipline.finance.execution import MarketOrder, LimitOrder
from zipline.finance.trading import TradingEnvironment
from zipline.finance.performance import PerformanceTracker
from zipline.finance.trading import SimulationParameters
from zipline.testing import (
setup_logger,
teardown_logger
)
from zipline.data.us_equity_pricing import BcolzDailyBarReader
from zipline.data.minute_bars import BcolzMinuteBarReader
from zipline.data.data_portal import DataPortal
from zipline.data.us_equity_pricing import BcolzDailyBarWriter
from zipline.finance.slippage import FixedSlippage
from zipline.protocol import BarData
from zipline.testing.core import write_bcolz_minute_data
from .utils.daily_bar_writer import DailyBarWriterFromDataFrames
from zipline.testing import (
tmp_trading_env,
write_bcolz_minute_data,
)
from zipline.testing.fixtures import (
WithLogger,
WithTradingEnvironment,
ZiplineTestCase,
)
import zipline.utils.factory as factory
@@ -54,26 +56,16 @@ EXTENDED_TIMEOUT = 90
_multiprocess_can_split_ = False
class FinanceTestCase(TestCase):
class FinanceTestCase(WithLogger,
WithTradingEnvironment,
ZiplineTestCase):
ASSET_FINDER_EQUITY_SIDS = 1, 2, 133
start = START_DATE = pd.Timestamp('2006-01-01', tz='utc')
end = END_DATE = pd.Timestamp('2006-12-31', tz='utc')
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.env.write_data(equities_identifiers=[1, 2, 133])
@classmethod
def tearDownClass(cls):
del cls.env
def setUp(self):
self.zipline_test_config = {
'sid': 133,
}
setup_logger(self)
def tearDown(self):
teardown_logger(self)
def init_instance_fixtures(self):
super(FinanceTestCase, self).init_instance_fixtures()
self.zipline_test_config = {'sid': 133}
# TODO: write tests for short sales
# TODO: write a test to do massive buying or shorting.
@@ -174,33 +166,35 @@ class FinanceTestCase(TestCase):
self.transaction_sim(**params1)
def transaction_sim(self, **params):
""" This is a utility method that asserts expected
"""This is a utility method that asserts expected
results for conversion of orders to transactions given a
trade history"""
tempdir = TempDirectory()
try:
trade_count = params['trade_count']
trade_interval = params['trade_interval']
order_count = params['order_count']
order_amount = params['order_amount']
order_interval = params['order_interval']
expected_txn_count = params['expected_txn_count']
expected_txn_volume = params['expected_txn_volume']
trade history
"""
trade_count = params['trade_count']
trade_interval = params['trade_interval']
order_count = params['order_count']
order_amount = params['order_amount']
order_interval = params['order_interval']
expected_txn_count = params['expected_txn_count']
expected_txn_volume = params['expected_txn_volume']
# optional parameters
# ---------------------
# if present, alternate between long and short sales
alternate = params.get('alternate')
# optional parameters
# ---------------------
# if present, alternate between long and short sales
alternate = params.get('alternate')
# if present, expect transaction amounts to match orders exactly.
complete_fill = params.get('complete_fill')
# if present, expect transaction amounts to match orders exactly.
complete_fill = params.get('complete_fill')
env = TradingEnvironment()
sid = 1
sid = 1
metadata = make_simple_equity_info([sid], self.start, self.end)
with TempDirectory() as tempdir, \
tmp_trading_env(equities=metadata) as env:
if trade_interval < timedelta(days=1):
sim_params = factory.create_simulation_parameters(
start=self.start,
end=self.end,
data_frequency="minute"
)
@@ -253,8 +247,7 @@ class FinanceTestCase(TestCase):
}
path = os.path.join(tempdir.path, "testdata.bcolz")
DailyBarWriterFromDataFrames(assets).write(
path, days, assets)
BcolzDailyBarWriter(path, days).write(assets.items())
equity_daily_reader = BcolzDailyBarReader(path)
@@ -272,13 +265,6 @@ class FinanceTestCase(TestCase):
blotter = Blotter(sim_params.data_frequency, self.env.asset_finder,
slippage_func)
env.write_data(equities_data={
sid: {
"start_date": sim_params.trading_days[0],
"end_date": sim_params.trading_days[-1]
}
})
start_date = sim_params.first_open
if alternate:
@@ -356,8 +342,6 @@ class FinanceTestCase(TestCase):
# the open orders should not contain sid.
oo = blotter.open_orders
self.assertNotIn(sid, oo, "Entry is removed when no open orders")
finally:
tempdir.cleanup()
def test_blotter_processes_splits(self):
blotter = Blotter('daily', self.env.asset_finder,
@@ -393,25 +377,12 @@ class FinanceTestCase(TestCase):
self.assertEqual(2, fls_order['sid'])
class TradingEnvironmentTestCase(TestCase):
class TradingEnvironmentTestCase(WithLogger,
WithTradingEnvironment,
ZiplineTestCase):
"""
Tests for date management utilities in zipline.finance.trading.
"""
def setUp(self):
setup_logger(self)
def tearDown(self):
teardown_logger(self)
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
@classmethod
def tearDownClass(cls):
del cls.env
@timed(DEFAULT_TIMEOUT)
def test_is_trading_day(self):
# holidays taken from: http://www.nyse.com/press/1191407641943.html
+373 -476
View File
File diff suppressed because it is too large Load Diff
+161 -158
View File
@@ -22,8 +22,6 @@ from datetime import (
)
import logging
from testfixtures import TempDirectory
import unittest
import nose.tools as nt
import pytz
@@ -32,6 +30,7 @@ import numpy as np
from six.moves import range, zip
from zipline.assets import Asset
from zipline.assets.synthetic import make_simple_equity_info
from zipline.data.us_equity_pricing import (
SQLiteAdjustmentWriter,
SQLiteAdjustmentReader,
@@ -43,14 +42,24 @@ import zipline.utils.math_utils as zp_math
from zipline.finance.blotter import Order
from zipline.finance.commission import PerShare, PerTrade, PerDollar
from zipline.finance.trading import TradingEnvironment
from zipline.finance.performance.position import Position
from zipline.utils.factory import create_simulation_parameters
from zipline.utils.serialization_utils import (
loads_with_persistent_ids, dumps_with_persistent_ids
)
from zipline.testing.core import create_data_portal_from_trade_history, \
create_empty_splits_mergers_frame
from zipline.testing import (
MockDailyBarReader,
create_data_portal_from_trade_history,
create_empty_splits_mergers_frame,
tmp_trading_env,
)
from zipline.testing.fixtures import (
WithInstanceTmpDir,
WithSimParams,
WithTmpDir,
WithTradingEnvironment,
ZiplineTestCase,
)
logger = logging.getLogger('Test Perf Tracking')
@@ -248,29 +257,25 @@ def setup_env_data(env, sim_params, sids, futures_sids=[]):
env.write_data(futures_data=futures_data)
class TestSplitPerformance(unittest.TestCase):
class TestSplitPerformance(WithSimParams, WithTmpDir, ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
END_DATE = pd.Timestamp('2006-01-04', tz='utc')
SIM_PARAMS_CAPITAL_BASE = 10e3
ASSET_FINDER_EQUITY_SIDS = 1, 2
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.sim_params = create_simulation_parameters(num_days=2,
capital_base=10e3)
setup_env_data(cls.env, cls.sim_params, [1, 2])
cls.tempdir = TempDirectory()
def init_class_fixtures(cls):
super(TestSplitPerformance, cls).init_class_fixtures()
cls.asset1 = cls.env.asset_finder.retrieve_asset(1)
@classmethod
def tearDownClass(cls):
cls.tempdir.cleanup()
def test_multiple_splits(self):
# if multiple positions all have splits at the same time, verify that
# the total leftover cash is correct
perf_tracker = perf.PerformanceTracker(self.sim_params, self.env)
asset1 = self.env.asset_finder.retrieve_asset(1)
asset2 = self.env.asset_finder.retrieve_asset(2)
asset1 = self.asset_finder.retrieve_asset(1)
asset2 = self.asset_finder.retrieve_asset(2)
perf_tracker.position_tracker.positions[1] = \
Position(asset1, amount=10, cost_basis=10, last_sale_price=11)
@@ -303,7 +308,7 @@ class TestSplitPerformance(unittest.TestCase):
# 100 shares at $20 apiece = $2000 position
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.tmpdir,
self.sim_params,
{1: events},
)
@@ -387,22 +392,17 @@ class TestSplitPerformance(unittest.TestCase):
(i, perf_kind, perf_result['returns']))
class TestCommissionEvents(unittest.TestCase):
class TestCommissionEvents(WithSimParams, WithTmpDir, ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
END_DATE = pd.Timestamp('2006-01-09', tz='utc')
ASSET_FINDER_EQUITY_SIDS = 0, 1, 133
SIM_PARAMS_CAPITAL_BASE = 10e3
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.sim_params = create_simulation_parameters(num_days=5,
capital_base=10e3)
setup_env_data(cls.env, cls.sim_params, [0, 1, 133])
cls.tempdir = TempDirectory()
def init_class_fixtures(cls):
super(TestCommissionEvents, cls).init_class_fixtures()
cls.asset1 = cls.env.asset_finder.retrieve_asset(1)
@classmethod
def tearDownClass(cls):
cls.tempdir.cleanup()
def test_commission_event(self):
trade_events = factory.create_trade_history(
self.asset1,
@@ -422,7 +422,7 @@ class TestCommissionEvents(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.tmpdir,
self.sim_params,
{1: trade_events},
)
@@ -506,7 +506,7 @@ class TestCommissionEvents(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.tmpdir,
self.sim_params,
{1: events},
)
@@ -548,7 +548,7 @@ class TestCommissionEvents(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.tmpdir,
self.sim_params,
{1: events},
)
@@ -569,34 +569,19 @@ class TestCommissionEvents(unittest.TestCase):
9700)
class MockDailyBarSpotReader(object):
def spot_price(self, sid, day, colname):
return 100.0
class TestDividendPerformance(unittest.TestCase):
class TestDividendPerformance(WithSimParams,
WithInstanceTmpDir,
ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
END_DATE = pd.Timestamp('2006-01-10', tz='utc')
ASSET_FINDER_EQUITY_SIDS = 1, 2
SIM_PARAMS_CAPITAL_BASE = 10e3
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.sim_params = create_simulation_parameters(num_days=6,
capital_base=10e3)
setup_env_data(cls.env, cls.sim_params, [1, 2])
cls.asset1 = cls.env.asset_finder.retrieve_asset(1)
cls.asset2 = cls.env.asset_finder.retrieve_asset(2)
@classmethod
def tearDownClass(cls):
del cls.env
def setUp(self):
self.tempdir = TempDirectory()
def tearDown(self):
self.tempdir.cleanup()
def init_class_fixtures(cls):
super(TestDividendPerformance, cls).init_class_fixtures()
cls.asset1 = cls.asset_finder.retrieve_asset(1)
cls.asset2 = cls.asset_finder.retrieve_asset(2)
def test_market_hours_calculations(self):
# DST in US/Eastern began on Sunday March 14, 2010
@@ -619,10 +604,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -634,10 +622,9 @@ class TestDividendPerformance(unittest.TestCase):
})
writer.write(splits, mergers, dividends)
adjustment_reader = SQLiteAdjustmentReader(dbpath)
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -682,10 +669,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([], dtype=np.uint32),
@@ -710,7 +700,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
events,
)
@@ -753,10 +743,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -771,7 +764,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -811,10 +804,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -829,7 +825,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -869,10 +865,13 @@ class TestDividendPerformance(unittest.TestCase):
self.sim_params,
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
@@ -888,7 +887,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -932,10 +931,13 @@ class TestDividendPerformance(unittest.TestCase):
for i in range(30):
pay_date = factory.get_next_trading_dt(pay_date, oneday, self.env)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -950,7 +952,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -990,10 +992,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -1008,7 +1013,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -1045,10 +1050,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -1063,7 +1071,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: events},
)
@@ -1098,10 +1106,13 @@ class TestDividendPerformance(unittest.TestCase):
env=self.env
)
dbpath = self.tempdir.getpath('adjustments.sqlite')
dbpath = self.instance_tmpdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath, self.env.trading_days,
MockDailyBarSpotReader())
writer = SQLiteAdjustmentWriter(
dbpath,
MockDailyBarReader(),
self.env.trading_days,
)
splits = mergers = create_empty_splits_mergers_frame()
dividends = pd.DataFrame({
'sid': np.array([1], dtype=np.uint32),
@@ -1128,7 +1139,7 @@ class TestDividendPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
sim_params,
{1: events},
)
@@ -1163,45 +1174,43 @@ class TestDividendPerformanceHolidayStyle(TestDividendPerformance):
# two days ahead. Any tests that hard code events
# to be start + oneday will fail, since those events will
# be skipped by the simulation.
START_DATE = pd.Timestamp('2003-11-30', tz='utc')
END_DATE = pd.Timestamp('2003-12-08', tz='utc')
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
cls.sim_params = create_simulation_parameters(
num_days=6,
capital_base=10e3,
start=pd.Timestamp("2003-11-30", tz='UTC'),
end=pd.Timestamp("2003-12-08", tz='UTC')
class TestPositionPerformance(WithInstanceTmpDir, ZiplineTestCase):
def create_environment_stuff(self,
num_days=4,
sids=[1, 2],
futures_sids=[3]):
start = pd.Timestamp('2006-01-01', tz='utc')
end = start + timedelta(days=num_days * 2)
equities = make_simple_equity_info(sids, start, end)
futures = pd.DataFrame.from_dict(
{
sid: {
'start_date': start,
'end_date': end,
'multiplier': 100,
}
for sid in futures_sids
},
orient='index',
)
self.env = self.enter_instance_context(tmp_trading_env(
equities=equities,
futures=futures,
))
self.sim_params = create_simulation_parameters(
start=start,
num_days=num_days,
)
setup_env_data(cls.env, cls.sim_params, [1, 2])
cls.asset1 = cls.env.asset_finder.retrieve_asset(1)
cls.asset2 = cls.env.asset_finder.retrieve_asset(2)
class TestPositionPerformance(unittest.TestCase):
def setUp(self):
self.tempdir = TempDirectory()
def create_environment_stuff(self, num_days=4, sids=[1, 2],
futures_sids=[3]):
self.env = TradingEnvironment()
self.sim_params = create_simulation_parameters(num_days=num_days)
setup_env_data(self.env, self.sim_params, sids, futures_sids)
self.finder = self.env.asset_finder
self.asset1 = self.env.asset_finder.retrieve_asset(1)
self.asset2 = self.env.asset_finder.retrieve_asset(2)
self.asset3 = self.env.asset_finder.retrieve_asset(3)
def tearDown(self):
self.tempdir.cleanup()
del self.env
def test_long_short_positions(self):
"""
start with $1000
@@ -1232,7 +1241,7 @@ class TestPositionPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: trades_1, 2: trades_2}
)
@@ -1328,7 +1337,7 @@ class TestPositionPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: trades})
txn = create_txn(self.asset1, trades[1].dt, 10.0, 1000)
@@ -1419,7 +1428,7 @@ class TestPositionPerformance(unittest.TestCase):
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: trades})
txn = create_txn(self.asset1, trades[1].dt, 10.0, 100)
@@ -1536,7 +1545,7 @@ single short-sale transaction"""
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: trades})
@@ -1767,7 +1776,7 @@ cost of sole txn in test"
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{3: trades}
)
@@ -1886,7 +1895,7 @@ single short-sale transaction"""
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{3: trades}
)
@@ -2130,7 +2139,7 @@ trade after cover"""
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: trades})
@@ -2218,7 +2227,7 @@ shares in position"
data_portal = create_data_portal_from_trade_history(
self.env,
self.tempdir,
self.instance_tmpdir,
self.sim_params,
{1: trades})
@@ -2363,27 +2372,21 @@ shares in position"
self.assertEqual(pp.positions[1].cost_basis, cost_bases[-1])
class TestPositionTracker(unittest.TestCase):
class TestPositionTracker(WithTradingEnvironment,
WithInstanceTmpDir,
ZiplineTestCase):
ASSET_FINDER_EQUITY_SIDS = 1, 2
@classmethod
def setUpClass(cls):
cls.env = TradingEnvironment()
futures_metadata = {3: {'multiplier': 1000},
4: {'multiplier': 1000},
1032201401: {'multiplier': 50},
}
cls.env.write_data(equities_identifiers=[1, 2],
futures_data=futures_metadata)
@classmethod
def tearDownClass(cls):
del cls.env
def setUp(self):
self.tempdir = TempDirectory()
def tearDown(self):
self.tempdir.cleanup()
def make_futures_info(cls):
return pd.DataFrame.from_dict(
{
3: {'multiplier': 1000},
4: {'multiplier': 1000},
1032201401: {'multiplier': 50},
},
orient='index',
)
def test_empty_positions(self):
"""
+55 -84
View File
@@ -1,19 +1,18 @@
import pandas as pd
from datetime import timedelta
from unittest import TestCase
import pandas as pd
from testfixtures import TempDirectory
from zipline.algorithm import TradingAlgorithm
from zipline.errors import TradingControlViolation
from zipline.finance.trading import TradingEnvironment
from zipline.testing import (
add_security_data,
create_data_portal,
security_list_copy,
setup_logger,
teardown_logger,
tmp_trading_env,
tmp_dir,
)
from zipline.testing.core import create_data_portal
from zipline.testing.fixtures import WithLogger, ZiplineTestCase
from zipline.utils import factory
from zipline.utils.security_list import (
SecurityListSet,
@@ -64,58 +63,47 @@ class IterateRLAlgo(TradingAlgorithm):
self.found = True
class SecurityListTestCase(TestCase):
class SecurityListTestCase(WithLogger, ZiplineTestCase):
@classmethod
def setUpClass(cls):
def init_class_fixtures(cls):
super(SecurityListTestCase, cls).init_class_fixtures()
# this is ugly, but we need to create two different
# TradingEnvironment/DataPortal pairs
cls.env = TradingEnvironment()
cls.env2 = TradingEnvironment()
cls.extra_knowledge_date = pd.Timestamp("2015-01-27", tz='UTC')
cls.trading_day_before_first_kd = pd.Timestamp("2015-01-23", tz='UTC')
start = list(LEVERAGED_ETFS.keys())[0]
end = pd.Timestamp('2015-02-17', tz='utc')
cls.extra_knowledge_date = pd.Timestamp('2015-01-27', tz='utc')
cls.trading_day_before_first_kd = pd.Timestamp('2015-01-23', tz='utc')
symbols = ['AAPL', 'GOOG', 'BZQ', 'URTY', 'JFT']
days = cls.env.days_in_range(
list(LEVERAGED_ETFS.keys())[0],
pd.Timestamp("2015-02-17", tz='UTC')
)
cls.env = cls.enter_class_context(tmp_trading_env(
equities=pd.DataFrame.from_records([{
'start_date': start,
'end_date': end,
'symbol': symbol
} for symbol in symbols]),
))
cls.sim_params = factory.create_simulation_parameters(
start=list(LEVERAGED_ETFS.keys())[0],
start=start,
num_days=4,
env=cls.env
)
cls.sim_params2 = factory.create_simulation_parameters(
cls.sim_params2 = sp2 = factory.create_simulation_parameters(
start=cls.trading_day_before_first_kd, num_days=4
)
equities_metadata = {}
for i, symbol in enumerate(symbols):
equities_metadata[i] = {
'start_date': days[0],
'end_date': days[-1],
cls.env2 = cls.enter_class_context(tmp_trading_env(
equities=pd.DataFrame.from_records([{
'start_date': sp2.period_start,
'end_date': sp2.period_end,
'symbol': symbol
}
} for symbol in symbols]),
))
equities_metadata2 = {}
for i, symbol in enumerate(symbols):
equities_metadata2[i] = {
'start_date': cls.sim_params2.period_start,
'end_date': cls.sim_params2.period_end,
'symbol': symbol
}
cls.env.write_data(equities_data=equities_metadata)
cls.env2.write_data(equities_data=equities_metadata2)
cls.tempdir = TempDirectory()
cls.tempdir2 = TempDirectory()
cls.tempdir = cls.enter_class_context(tmp_dir())
cls.tempdir2 = cls.enter_class_context(tmp_dir())
cls.data_portal = create_data_portal(
env=cls.env,
@@ -131,15 +119,6 @@ class SecurityListTestCase(TestCase):
sids=range(0, 5)
)
setup_logger(cls)
@classmethod
def tearDownClass(cls):
del cls.env
cls.tempdir.cleanup()
cls.tempdir2.cleanup()
teardown_logger(cls)
def test_iterate_over_restricted_list(self):
algo = IterateRLAlgo(symbol='BZQ', sim_params=self.sim_params,
env=self.env)
@@ -271,41 +250,33 @@ class SecurityListTestCase(TestCase):
self.check_algo_exception(algo, ctx, 0)
def test_algo_without_rl_violation_after_delete(self):
new_tempdir = TempDirectory()
try:
with security_list_copy():
# add a delete statement removing bzq
# write a new delete statement file to disk
add_security_data([], ['BZQ'])
sim_params = factory.create_simulation_parameters(
start=self.extra_knowledge_date,
num_days=4,
)
equities = pd.DataFrame.from_records([{
'symbol': 'BZQ',
'start_date': sim_params.period_start,
'end_date': sim_params.period_end,
}])
with TempDirectory() as new_tempdir, \
security_list_copy(), \
tmp_trading_env(equities=equities) as env:
# add a delete statement removing bzq
# write a new delete statement file to disk
add_security_data([], ['BZQ'])
# now fast-forward to self.extra_knowledge_date. requires
# a new env, simparams, and dataportal
env = TradingEnvironment()
sim_params = factory.create_simulation_parameters(
start=self.extra_knowledge_date, num_days=4, env=env)
data_portal = create_data_portal(
env,
new_tempdir,
sim_params,
range(0, 5)
)
env.write_data(equities_data={
"0": {
'symbol': 'BZQ',
'start_date': sim_params.period_start,
'end_date': sim_params.period_end,
}
})
data_portal = create_data_portal(
env,
new_tempdir,
sim_params,
range(0, 5)
)
algo = RestrictedAlgoWithoutCheck(
symbol='BZQ', sim_params=sim_params, env=env
)
algo.run(data_portal)
finally:
new_tempdir.cleanup()
algo = RestrictedAlgoWithoutCheck(
symbol='BZQ', sim_params=sim_params, env=env
)
algo.run(data_portal)
def test_algo_with_rl_violation_after_add(self):
with security_list_copy():
-60
View File
@@ -1,60 +0,0 @@
#
# Copyright 2013 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import TestCase
import pandas as pd
import pytz
import numpy as np
from zipline.utils.factory import (load_from_yahoo,
load_bars_from_yahoo)
class TestFactory(TestCase):
def test_load_from_yahoo(self):
stocks = ['AAPL', 'GE']
start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)
end = pd.datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc)
data = load_from_yahoo(stocks=stocks, start=start, end=end)
assert data.index[0] == pd.Timestamp('1993-01-04 00:00:00+0000')
assert data.index[-1] == pd.Timestamp('2001-12-31 00:00:00+0000')
for stock in stocks:
assert stock in data.columns
np.testing.assert_raises(
AssertionError, load_from_yahoo, stocks=stocks,
start=end, end=start
)
def test_load_bars_from_yahoo(self):
stocks = ['AAPL', 'GE']
start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)
end = pd.datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc)
data = load_bars_from_yahoo(stocks=stocks, start=start, end=end)
assert data.major_axis[0] == pd.Timestamp('1993-01-04 00:00:00+0000')
assert data.major_axis[-1] == pd.Timestamp('2001-12-31 00:00:00+0000')
for stock in stocks:
assert stock in data.items
for ohlc in ['open', 'high', 'low', 'close', 'volume', 'price']:
assert ohlc in data.minor_axis
np.testing.assert_raises(
AssertionError, load_bars_from_yahoo, stocks=stocks,
start=end, end=start
)
+2 -2
View File
@@ -5,9 +5,9 @@ from six import with_metaclass
from zipline.utils.final import (
FinalMeta,
final_meta_factory,
final,
)
from zipline.utils.metautils import compose_types
class FinalMetaTestCase(TestCase):
@@ -159,7 +159,7 @@ class FinalMetaTestCase(TestCase):
class FinalABCMetaTestCase(FinalMetaTestCase):
@classmethod
def setUpClass(cls):
FinalABCMeta = final_meta_factory(ABCMeta)
FinalABCMeta = compose_types(FinalMeta, ABCMeta)
class ABCWithFinal(with_metaclass(FinalABCMeta, object)):
a = final('ABCWithFinal: a')
+52 -21
View File
@@ -12,8 +12,9 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from copy import copy
import operator as op
import warnings
import logbook
import pytz
@@ -33,6 +34,7 @@ from six import (
)
from zipline._protocol import handle_non_market_minutes
from zipline.assets.synthetic import make_simple_equity_info
from zipline.data.data_portal import DataPortal
from zipline.errors import (
AttachPipelineAfterInitialize,
@@ -96,6 +98,7 @@ from zipline.utils.events import (
TimeRuleFactory,
)
from zipline.utils.factory import create_simulation_parameters
from zipline.utils.functional import unzip
from zipline.utils.math_utils import (
tolerant_equals,
round_if_near_integer
@@ -252,11 +255,18 @@ class TradingAlgorithm(object):
self.trading_environment = TradingEnvironment()
# Update the TradingEnvironment with the provided asset metadata
self.trading_environment.write_data(
equities_data=kwargs.pop('equities_metadata', {}),
equities_identifiers=kwargs.pop('identifiers', []),
futures_data=kwargs.pop('futures_metadata', {}),
)
if 'equities_metadata' in kwargs or 'futures_metadata' in kwargs:
warnings.warn(
'passing metadata to TradingAlgorithm is deprecated; please'
' write this data into the asset db before passing it to the'
' trading environment',
DeprecationWarning,
stacklevel=1,
)
self.trading_environment.write_data(
equities=kwargs.pop('equities_metadata', None),
futures=kwargs.pop('futures_metadata', None),
)
# set the capital base
self.capital_base = kwargs.pop('capital_base', DEFAULT_CAPITAL_BASE)
@@ -563,6 +573,17 @@ class TradingAlgorithm(object):
data = data.swapaxes(0, 2)
if isinstance(data, pd.Panel):
# For compatibility with existing examples allow start/end
# to be inferred.
if overwrite_sim_params:
self.sim_params.period_start = data.major_axis[0]
self.sim_params.period_end = data.major_axis[-1]
# Changing period_start and period_close might require
# updating of first_open and last_close.
self.sim_params.update_internal_from_env(
env=self.trading_environment
)
copy_panel = data.copy()
copy_panel.items = self._write_and_map_id_index_to_sids(
copy_panel.items, copy_panel.major_axis[0],
@@ -586,17 +607,6 @@ class TradingAlgorithm(object):
self.trading_environment,
equity_daily_reader=equity_daily_reader)
# For compatibility with existing examples allow start/end
# to be inferred.
if overwrite_sim_params:
self.sim_params.period_start = data.major_axis[0]
self.sim_params.period_end = data.major_axis[-1]
# Changing period_start and period_close might require
# updating of first_open and last_close.
self.sim_params.update_internal_from_env(
env=self.trading_environment
)
# Force a reset of the performance tracker, in case
# this is a repeat run of the algorithm.
self.perf_tracker = None
@@ -620,7 +630,8 @@ class TradingAlgorithm(object):
def _write_and_map_id_index_to_sids(self, identifiers, as_of_date):
# Build new Assets for identifiers that can't be resolved as
# sids/Assets
identifiers_to_build = []
identifiers_to_build = set()
next_sid = max(self.asset_finder.sids or (0,)) + 1
for identifier in identifiers:
asset = None
@@ -631,10 +642,30 @@ class TradingAlgorithm(object):
asset = self.asset_finder.retrieve_asset(sid=identifier,
default_none=True)
if asset is None:
identifiers_to_build.append(identifier)
try:
sid = op.index(identifier)
except TypeError:
sid = next_sid
next_sid += 1
identifiers_to_build.add((identifier, sid))
self.trading_environment.write_data(
equities_identifiers=identifiers_to_build)
if identifiers_to_build:
warnings.warn(
'writing unknown identifiers into the assets db of the trading'
' environment is deprecated; please write this information'
' to the assets db before constructing the environment',
DeprecationWarning,
stacklevel=2,
)
symbols, sids = unzip(identifiers_to_build, 2)
self.trading_environment.write_data(
equities=make_simple_equity_info(
sids,
start_date=self.sim_params.period_start,
end_date=self.sim_params.period_end,
symbols=symbols,
),
)
# We need to clear out any cache misses that were stored while trying
# to do lookups. The real fix for this problem is to not construct an
+4
View File
@@ -25,9 +25,13 @@ from .assets import (
AssetConvertible,
AssetFinderCachedEquities
)
from .asset_db_schema import ASSET_DB_VERSION
from .asset_writer import AssetDBWriter
__all__ = [
'ASSET_DB_VERSION',
'Asset',
'AssetDBWriter',
'Equity',
'Future',
'AssetFinder',
+88 -17
View File
@@ -1,6 +1,9 @@
import sqlalchemy as sa
from functools import wraps
from alembic.migration import MigrationContext
from alembic.operations import Operations
import sqlalchemy as sa
from toolz.curried import do, operator as op
from zipline.assets.asset_writer import write_version_info
from zipline.errors import AssetDBImpossibleDowngrade
@@ -68,13 +71,46 @@ def _pragma_foreign_keys(connection, on):
connection.execute("PRAGMA foreign_keys=%s" % ("ON" if on else "OFF"))
def _downgrade_v1_to_v0(op, version_info_table):
# This dict contains references to downgrade methods that can be applied to an
# assets db. The resulting db's version is the key.
# e.g. The method at key '0' is the downgrade method from v1 to v0
_downgrade_methods = {}
def downgrades(src):
"""Decorator for marking that a method is a downgrade to a version to the
previous version.
Parameters
----------
src : int
The version this downgrades from.
Returns
-------
decorator : callable[(callable) -> callable]
The decorator to apply.
"""
def _(f):
destination = src - 1
@do(op.setitem(_downgrade_methods, destination))
@wraps(f)
def wrapper(op, version_info_table):
version_info_table.delete().execute() # clear the version
f(op)
write_version_info(version_info_table, destination)
return wrapper
return _
@downgrades(1)
def _downgrade_v1(op):
"""
Downgrade assets db by removing the 'tick_size' column and renaming the
'multiplier' column.
"""
version_info_table.delete().execute()
# Drop indices before batch
# This is to prevent index collision when creating the temp table
op.drop_index('ix_futures_contracts_root_symbol')
@@ -99,15 +135,12 @@ def _downgrade_v1_to_v0(op, version_info_table):
columns=['symbol'],
unique=True)
write_version_info(version_info_table, 0)
def _downgrade_v2_to_v1(op, version_info_table):
@downgrades(2)
def _downgrade_v2(op):
"""
Downgrade assets db by removing the 'auto_close_date' column.
"""
version_info_table.delete().execute()
# Drop indices before batch
# This is to prevent index collision when creating the temp table
op.drop_index('ix_equities_fuzzy_symbol')
@@ -126,12 +159,50 @@ def _downgrade_v2_to_v1(op, version_info_table):
table_name='equities',
columns=['company_symbol'])
write_version_info(version_info_table, 1)
# This dict contains references to downgrade methods that can be applied to an
# assets db. The resulting db's version is the key.
# e.g. The method at key '0' is the downgrade method from v1 to v0
_downgrade_methods = {
0: _downgrade_v1_to_v0,
1: _downgrade_v2_to_v1,
}
@downgrades(3)
def _downgrade_v3(op):
"""
Downgrade assets db by adding a not null constraint on
``equities.first_traded``
"""
op.create_table(
'_new_equities',
sa.Column(
'sid',
sa.Integer,
unique=True,
nullable=False,
primary_key=True,
),
sa.Column('symbol', sa.Text),
sa.Column('company_symbol', sa.Text),
sa.Column('share_class_symbol', sa.Text),
sa.Column('fuzzy_symbol', sa.Text),
sa.Column('asset_name', sa.Text),
sa.Column('start_date', sa.Integer, default=0, nullable=False),
sa.Column('end_date', sa.Integer, nullable=False),
sa.Column('first_traded', sa.Integer, nullable=False),
sa.Column('auto_close_date', sa.Integer),
sa.Column('exchange', sa.Text),
)
op.execute(
"""
insert into _new_equities
select * from equities
where equities.first_traded is not null
""",
)
op.drop_table('equities')
op.rename_table('_new_equities', 'equities')
# we need to make sure the indicies have the proper names after the rename
op.create_index(
'ix_equities_company_symbol',
'equities',
['company_symbol'],
)
op.create_index(
'ix_equities_fuzzy_symbol',
'equities',
['fuzzy_symbol'],
)
+14 -7
View File
@@ -4,7 +4,9 @@ import sqlalchemy as sa
# Define a version number for the database generated by these writers
# Increment this version number any time a change is made to the schema of the
# assets database
ASSET_DB_VERSION = 2
# NOTE: When upgrading this remember to add a downgrade in:
# .asset_db_migrations
ASSET_DB_VERSION = 3
def generate_asset_db_metadata(bind=None):
@@ -19,11 +21,16 @@ def generate_asset_db_metadata(bind=None):
return metadata
# A list of the names of all tables in the assets db
# A frozenset of the names of all tables in the assets db
# NOTE: When modifying this schema, update the ASSET_DB_VERSION value
asset_db_table_names = ['version_info', 'equities', 'futures_exchanges',
'futures_root_symbols', 'futures_contracts',
'asset_router']
asset_db_table_names = frozenset({
'asset_router',
'equities',
'futures_contracts',
'futures_exchanges',
'futures_root_symbols',
'version_info',
})
def _equities_table_schema(metadata):
@@ -45,7 +52,7 @@ def _equities_table_schema(metadata):
sa.Column('asset_name', sa.Text),
sa.Column('start_date', sa.Integer, default=0, nullable=False),
sa.Column('end_date', sa.Integer, nullable=False),
sa.Column('first_traded', sa.Integer, nullable=False),
sa.Column('first_traded', sa.Integer),
sa.Column('auto_close_date', sa.Integer),
sa.Column('exchange', sa.Text),
)
@@ -112,7 +119,7 @@ def _futures_contracts_schema(metadata):
sa.Column('asset_name', sa.Text),
sa.Column('start_date', sa.Integer, default=0, nullable=False),
sa.Column('end_date', sa.Integer, nullable=False),
sa.Column('first_traded', sa.Integer, nullable=False),
sa.Column('first_traded', sa.Integer),
sa.Column(
'exchange',
sa.Text,
+201 -322
View File
@@ -12,32 +12,27 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import (
ABCMeta,
abstractmethod,
)
from collections import namedtuple
import re
import pandas as pd
import numpy as np
from six import with_metaclass
import sqlalchemy as sa
from zipline.errors import SidAssignmentError, AssetDBVersionError
from zipline.assets._assets import Asset
from contextlib2 import ExitStack
import numpy as np
import pandas as pd
import sqlalchemy as sa
from toolz import first
from zipline.errors import AssetDBVersionError
from zipline.assets.asset_db_schema import (
generate_asset_db_metadata,
asset_db_table_names,
ASSET_DB_VERSION,
)
SQLITE_MAX_VARIABLE_NUMBER = 999
# Define a namedtuple for use with the load_data and _load_data methods
AssetData = namedtuple('AssetData', 'equities futures exchanges root_symbols')
SQLITE_MAX_VARIABLE_NUMBER = 999
# Default values for the equities DataFrame
_equities_defaults = {
'symbol': None,
@@ -171,6 +166,27 @@ def _generate_output_dataframe(data_subset, defaults):
return output
def _dt_to_epoch_ns(dt_series):
"""Convert a timeseries into an Int64Index of nanoseconds since the epoch.
Parameters
----------
dt_series : pd.Series
The timeseries to convert.
Returns
-------
idx : pd.Int64Index
The index converted to nanoseconds since the epoch.
"""
index = pd.to_datetime(dt_series.values)
if index.tzinfo is None:
index = index.tz_localize('UTC')
else:
index = index.tz_convert('UTC')
return index.view(np.int64)
def check_version_info(version_table, expected_version):
"""
Checks for a version value in the version table.
@@ -216,171 +232,168 @@ def write_version_info(version_table, version_value):
sa.insert(version_table, values={'version': version_value}).execute()
class AssetDBWriter(with_metaclass(ABCMeta)):
class _empty(object):
columns = ()
class AssetDBWriter(object):
"""Class used to write data to an assets db.
Parameters
----------
engine : Engine or str
An SQLAlchemy engine or path to a SQL database.
"""
Class used to write arbitrary data to SQLite database.
Concrete subclasses will implement the logic for a specific
input datatypes by implementing the _load_data method.
DEFAULT_CHUNK_SIZE = SQLITE_MAX_VARIABLE_NUMBER
Methods
-------
write_all(engine, allow_sid_assignment=True, constraints=False)
Write the data supplied at initialization to the database.
init_db(engine, constraints=False)
Create the SQLite tables (called by write_all).
load_data()
Returns data in standard format.
def __init__(self, engine):
if isinstance(engine, str):
engine = sa.create_engine('sqlite:///' + engine)
self.engine = engine
"""
CHUNK_SIZE = SQLITE_MAX_VARIABLE_NUMBER
def write(self,
equities=None,
futures=None,
exchanges=None,
root_symbols=None,
chunk_size=DEFAULT_CHUNK_SIZE):
def __init__(self, equities=None, futures=None, exchanges=None,
root_symbols=None):
with self.engine.begin() as txn:
# Create SQL tables if they do not exist.
metadata = self.init_db(txn)
if equities is None:
equities = self.defaultval()
self._equities = equities
if futures is None:
futures = self.defaultval()
self._futures = futures
if exchanges is None:
exchanges = self.defaultval()
self._exchanges = exchanges
if root_symbols is None:
root_symbols = self.defaultval()
self._root_symbols = root_symbols
@abstractmethod
def defaultval(self):
raise NotImplementedError
def write_all(self,
engine,
allow_sid_assignment=True):
""" Write pre-supplied data to SQLite.
Parameters
----------
engine : Engine
An SQLAlchemy engine to a SQL database.
allow_sid_assignment: bool, optional
If True then the class can assign sids where necessary.
constraints : bool, optional
If True then create SQL ForeignKey and PrimaryKey constraints.
"""
self.allow_sid_assignment = allow_sid_assignment
# Begin an SQL transaction.
with engine.begin() as txn:
# Create SQL tables.
self.init_db(txn)
# Get the data to add to SQL.
data = self.load_data()
# Write the data to SQL.
self._write_exchanges(data.exchanges, txn)
self._write_root_symbols(data.root_symbols, txn)
self._write_futures(data.futures, txn)
self._write_equities(data.equities, txn)
data = self._load_data(
equities if equities is not None else pd.DataFrame(),
futures if futures is not None else pd.DataFrame(),
exchanges if exchanges is not None else pd.DataFrame(),
root_symbols if root_symbols is not None else pd.DataFrame(),
)
def _write_df_to_table(self, df, tbl, bind):
# Write the data to SQL.
self._write_df_to_table(
metadata.tables['futures_exchanges'],
data.exchanges,
txn,
chunk_size,
)
self._write_df_to_table(
metadata.tables['futures_root_symbols'],
data.root_symbols,
txn,
chunk_size,
)
asset_router = metadata.tables['asset_router']
self._write_assets(
asset_router,
metadata.tables['futures_contracts'],
'future',
data.futures,
txn,
chunk_size,
)
self._write_assets(
asset_router,
metadata.tables['equities'],
'equity',
data.equities,
txn,
chunk_size,
)
def _write_df_to_table(self, tbl, df, txn, chunk_size):
df.to_sql(
tbl.name,
bind.connection,
index_label=[col.name for col in tbl.primary_key.columns][0],
txn.connection,
index_label=first(tbl.primary_key.columns).name,
if_exists='append',
chunksize=self.CHUNK_SIZE,
chunksize=chunk_size,
)
def _write_assets(self, assets, asset_tbl, asset_type, bind):
self._write_df_to_table(assets, asset_tbl, bind)
def _write_assets(self,
asset_router,
tbl,
asset_type,
assets,
txn,
chunk_size):
self._write_df_to_table(tbl, assets, txn, chunk_size)
pd.DataFrame({self.asset_router.c.sid.name: assets.index.values,
self.asset_router.c.asset_type.name: asset_type}).to_sql(
self.asset_router.name,
bind.connection,
pd.DataFrame({
asset_router.c.sid.name: assets.index.values,
asset_router.c.asset_type.name: asset_type,
}).to_sql(
asset_router.name,
txn.connection,
if_exists='append',
index=False,
chunksize=self.CHUNK_SIZE,
chunksize=chunk_size
)
def _write_exchanges(self, exchanges, bind):
self._write_df_to_table(exchanges, self.futures_exchanges, bind)
def _write_root_symbols(self, root_symbols, bind):
self._write_df_to_table(root_symbols, self.futures_root_symbols, bind)
def _write_futures(self, futures, bind):
self._write_assets(futures, self.futures_contracts, 'future', bind)
def _write_equities(self, equities, bind):
self._write_assets(equities, self.equities, 'equity', bind)
def check_for_tables(self, engine):
def _all_tables_present(self, txn):
"""
Checks if any tables are present in the current assets database.
Parameters
----------
txn : Transaction
The open transaction to check in.
Returns
-------
bool
has_tables : bool
True if any tables are present, otherwise False.
"""
conn = engine.connect()
conn = txn.connect()
for table_name in asset_db_table_names:
if engine.dialect.has_table(conn, table_name):
if txn.dialect.has_table(conn, table_name):
return True
return False
def init_db(self, engine):
def init_db(self, txn=None):
"""Connect to database and create tables.
Parameters
----------
engine : Engine
An engine to a SQL database.
constraints : bool, optional
If True, create SQL ForeignKey and PrimaryKey constraints.
txn : sa.engine.Connection, optional
The transaction to execute in. If this is not provided, a new
transaction will be started with the engine provided.
Returns
-------
metadata : sa.MetaData
The metadata that describes the new assets db.
"""
tables_already_exist = self.check_for_tables(engine)
metadata = generate_asset_db_metadata(bind=engine)
with ExitStack() as stack:
if txn is None:
txn = stack.enter_context(self.engine.begin())
for table_name in asset_db_table_names:
setattr(self, table_name, metadata.tables[table_name])
tables_already_exist = self._all_tables_present(txn)
metadata = generate_asset_db_metadata(bind=txn)
# Create the SQL tables if they do not already exist.
metadata.create_all(checkfirst=True)
# Create the SQL tables if they do not already exist.
metadata.create_all(checkfirst=True)
if tables_already_exist:
check_version_info(self.version_info, ASSET_DB_VERSION)
else:
write_version_info(self.version_info, ASSET_DB_VERSION)
version_info = metadata.tables['version_info']
if tables_already_exist:
check_version_info(version_info, ASSET_DB_VERSION)
else:
write_version_info(version_info, ASSET_DB_VERSION)
return metadata
return metadata
def load_data(self):
"""
Returns a standard set of pandas.DataFrames:
equities, futures, exchanges, root_symbols
"""
def _normalize_equities(self, equities):
# HACK: If 'company_name' is provided, map it to asset_name
if ('company_name' in equities.columns and
'asset_name' not in equities.columns):
equities['asset_name'] = equities['company_name']
data = self._load_data()
###############################
# Generate equities DataFrame #
###############################
# HACK: If company_name is provided, map it to asset_name
if ('company_name' in data.equities.columns
and 'asset_name' not in data.equities.columns):
data.equities['asset_name'] = data.equities['company_name']
if 'file_name' in data.equities.columns:
data.equities['symbol'] = data.equities['file_name']
# remap 'file_name' to 'symbol' if provided
if 'file_name' in equities.columns:
equities['symbol'] = equities['file_name']
equities_output = _generate_output_dataframe(
data_subset=data.equities,
data_subset=equities,
defaults=_equities_defaults,
)
@@ -394,204 +407,70 @@ class AssetDBWriter(with_metaclass(ABCMeta)):
equities_output = equities_output.join(split_symbols)
# Upper-case all symbol data
equities_output['symbol'] = \
equities_output.symbol.str.upper()
equities_output['company_symbol'] = \
equities_output.company_symbol.str.upper()
equities_output['share_class_symbol'] = \
equities_output.share_class_symbol.str.upper()
equities_output['fuzzy_symbol'] = \
equities_output.fuzzy_symbol.str.upper()
for col in ('symbol',
'company_symbol',
'share_class_symbol',
'fuzzy_symbol'):
equities_output[col] = equities_output[col].str.upper()
# Convert date columns to UNIX Epoch integers (nanoseconds)
for date_col in ('start_date', 'end_date', 'first_traded',
'auto_close_date'):
equities_output[date_col] = \
self.dt_to_epoch_ns(equities_output[date_col])
for col in ('start_date',
'end_date',
'first_traded',
'auto_close_date'):
equities_output[col] = _dt_to_epoch_ns(equities_output[col])
##############################
# Generate futures DataFrame #
##############################
return equities_output
def _normalize_futures(self, futures):
futures_output = _generate_output_dataframe(
data_subset=data.futures,
data_subset=futures,
defaults=_futures_defaults,
)
for col in ('symbol', 'root_symbol'):
futures_output[col] = futures_output[col].str.upper()
# Convert date columns to UNIX Epoch integers (nanoseconds)
for date_col in ('start_date', 'end_date', 'first_traded',
'notice_date', 'expiration_date', 'auto_close_date'):
futures_output[date_col] = \
self.dt_to_epoch_ns(futures_output[date_col])
for col in ('start_date',
'end_date',
'first_traded',
'notice_date',
'expiration_date',
'auto_close_date'):
futures_output[col] = _dt_to_epoch_ns(futures_output[col])
# Convert symbols and root_symbols to upper case.
futures_output['symbol'] = futures_output.symbol.str.upper()
futures_output['root_symbol'] = futures_output.root_symbol.str.upper()
return futures_output
################################
# Generate exchanges DataFrame #
################################
exchanges_output = _generate_output_dataframe(
data_subset=data.exchanges,
defaults=_exchanges_defaults,
)
###################################
# Generate root symbols DataFrame #
###################################
root_symbols_output = _generate_output_dataframe(
data_subset=data.root_symbols,
defaults=_root_symbols_defaults,
)
return AssetData(equities=equities_output,
futures=futures_output,
exchanges=exchanges_output,
root_symbols=root_symbols_output)
@staticmethod
def dt_to_epoch_ns(dt_series):
index = pd.to_datetime(dt_series.values)
try:
index = index.tz_localize('UTC')
except TypeError:
index = index.tz_convert('UTC')
return index.view(np.int64)
@abstractmethod
def _load_data(self):
def _load_data(self, equities, futures, exchanges, root_symbols):
"""
Subclasses should implement this method to return data in a standard
format: a pandas.DataFrame for each of the following tables:
equities, futures, exchanges, root_symbols.
For each of these DataFrames the index columns should be the integer
unique identifier for the table, which are sid, sid, exchange_id and
root_symbol_id respectively.
Returns a standard set of pandas.DataFrames:
equities, futures, exchanges, root_symbols
"""
raise NotImplementedError('load_data')
class AssetDBWriterFromList(AssetDBWriter):
"""
Class used to write list data to SQLite database.
"""
defaultval = list
def _load_data(self):
# 0) Instantiate empty dictionaries
_equities, _futures, _exchanges, _root_symbols = {}, {}, {}, {}
# 1) Populate dictionaries
# Return the largest sid in our database, if one exists.
id_counter = sa.select(
[sa.func.max(self.asset_router.c.sid)]
).execute().scalar()
# Base sid creation on largest sid in database, or 0 if
# no sids exist.
if id_counter is None:
id_counter = 0
else:
id_counter += 1
for output, data in [(_equities, self._equities),
(_futures, self._futures), ]:
for identifier in data:
if isinstance(identifier, Asset):
sid = identifier.sid
metadata = identifier.to_dict()
output[sid] = metadata
elif hasattr(identifier, '__int__'):
output[identifier.__int__()] = {'symbol': None}
else:
if self.allow_sid_assignment:
output[id_counter] = {'symbol': identifier}
id_counter += 1
else:
raise SidAssignmentError(identifier=identifier)
exchange_counter = 0
for identifier in self._exchanges:
if hasattr(identifier, '__int__'):
_exchanges[identifier.__int__()] = {}
else:
_exchanges[exchange_counter] = {'exchange': identifier}
exchange_counter += 1
root_symbol_counter = 0
for identifier in self._root_symbols:
if hasattr(identifier, '__int__'):
_root_symbols[identifier.__int__()] = {}
else:
_root_symbols[root_symbol_counter] = \
{'root_symbol': identifier}
root_symbol_counter += 1
# 2) Convert dictionaries to pandas.DataFrames.
_equities = pd.DataFrame.from_dict(_equities, orient='index')
_futures = pd.DataFrame.from_dict(_futures, orient='index')
_exchanges = pd.DataFrame.from_dict(_exchanges, orient='index')
_root_symbols = pd.DataFrame.from_dict(_root_symbols, orient='index')
# 3) Return the data inside a named tuple.
return AssetData(equities=_equities,
futures=_futures,
exchanges=_exchanges,
root_symbols=_root_symbols)
class AssetDBWriterFromDictionary(AssetDBWriter):
"""
Class used to write dictionary data to SQLite database.
Expects to be initialised with dictionaries in the following format:
{id_0: {attribute_1 : ...}, id_1: {attribute_2: ...}, ...}
"""
defaultval = dict
def _load_data(self):
_equities = pd.DataFrame.from_dict(self._equities, orient='index')
_futures = pd.DataFrame.from_dict(self._futures, orient='index')
_exchanges = pd.DataFrame.from_dict(self._exchanges, orient='index')
_root_symbols = pd.DataFrame.from_dict(self._root_symbols,
orient='index')
return AssetData(equities=_equities,
futures=_futures,
exchanges=_exchanges,
root_symbols=_root_symbols)
class AssetDBWriterFromDataFrame(AssetDBWriter):
"""
Class used to write pandas.DataFrame data to SQLite database.
"""
defaultval = pd.DataFrame
def _load_data(self):
# Check whether identifier columns have been provided.
# If they have, set the index to this column.
# If not, assume the index already cotains the identifier information.
for df, id_col in [
(self._equities, 'sid'),
(self._futures, 'sid'),
(self._exchanges, 'exchange'),
(self._root_symbols, 'root_symbol'),
]:
for df, id_col in [(equities, 'sid'),
(futures, 'sid'),
(exchanges, 'exchange'),
(root_symbols, 'root_symbol')]:
if id_col in df.columns:
df.set_index([id_col], inplace=True)
df.set_index(id_col, inplace=True)
return AssetData(equities=self._equities,
futures=self._futures,
exchanges=self._exchanges,
root_symbols=self._root_symbols)
equities_output = self._normalize_equities(equities)
futures_output = self._normalize_futures(futures)
exchanges_output = _generate_output_dataframe(
data_subset=exchanges,
defaults=_exchanges_defaults,
)
root_symbols_output = _generate_output_dataframe(
data_subset=root_symbols,
defaults=_root_symbols_defaults,
)
return AssetData(
equities=equities_output,
futures=futures_output,
exchanges=exchanges_output,
root_symbols=root_symbols_output,
)
+47 -26
View File
@@ -101,6 +101,8 @@ class AssetFinder(object):
PERSISTENT_TOKEN = "<AssetFinder>"
def __init__(self, engine):
if isinstance(engine, str):
engine = sa.create_engine('sqlite:///' + engine)
self.engine = engine
metadata = sa.MetaData(bind=engine)
@@ -212,7 +214,7 @@ class AssetFinder(object):
Returns
-------
assets : list[int or None]
assets : list[Asset or None]
A list of the same length as `sids` containing Assets (or Nones)
corresponding to the requested sids.
@@ -684,12 +686,30 @@ class AssetFinder(object):
return sids
@property
def sids(self):
return tuple(map(
itemgetter('sid'),
sa.select((self.asset_router.c.sid,)).execute().fetchall(),
))
def _make_sids(tblattr):
def _(self):
return tuple(map(
itemgetter('sid'),
sa.select((
getattr(self, tblattr).c.sid,
)).execute().fetchall(),
))
return _
sids = property(
_make_sids('asset_router'),
doc='All the sids in the asset finder.',
)
equities_sids = property(
_make_sids('equities'),
doc='All of the sids for equities in the asset finder.',
)
futures_sids = property(
_make_sids('futures_contracts'),
doc='All of the sids for futures consracts in the asset finder.',
)
del _make_sids
def _lookup_generic_scalar(self,
asset_convertible,
@@ -928,35 +948,35 @@ class NotAssetConvertible(ValueError):
class AssetFinderCachedEquities(AssetFinder):
"""
An extension to AssetFinder that loads all equities from equities table
into memory and overrides the methods that lookup_symbol uses to look up
those equities.
An extension to AssetFinder that preloads all equities from equities table
into memory and does lookups from there.
To have any changes in the underlying assets db reflected by this asset
finder one must manually call the ``rehash_equities`` method.
"""
def __init__(self, engine):
super(AssetFinderCachedEquities, self).__init__(engine)
self.fuzzy_symbol_hashed_equities = {}
self.company_share_class_hashed_equities = {}
self.hashed_equities = sa.select(self.equities.c).execute().fetchall()
self._load_hashed_equities()
self._fuzzy_symbol_cache = {}
self._company_share_class_cache = {}
def _load_hashed_equities(self):
self.rehash_equities()
def rehash_equities(self):
"""Reload the underlying assets db into the in memory cache.
"""
Populates two maps - fuzzy symbol to list of equities having that
fuzzy symbol and company symbol/share class symbol to list of
equities having that combination of company symbol/share class symbol.
"""
for equity in self.hashed_equities:
for equity in sa.select(self.equities.c).execute().fetchall():
company_symbol = equity['company_symbol']
share_class_symbol = equity['share_class_symbol']
fuzzy_symbol = equity['fuzzy_symbol']
asset = self._convert_row_to_equity(equity)
self.company_share_class_hashed_equities.setdefault(
self._company_share_class_cache.setdefault(
(company_symbol, share_class_symbol),
[]
).append(asset)
self.fuzzy_symbol_hashed_equities.setdefault(
fuzzy_symbol, []
self._fuzzy_symbol_cache.setdefault(
fuzzy_symbol,
[],
).append(asset)
def _convert_row_to_equity(self, row):
@@ -966,7 +986,7 @@ class AssetFinderCachedEquities(AssetFinder):
return Equity(**_convert_asset_timestamp_fields(dict(row)))
def _get_fuzzy_candidates(self, fuzzy_symbol):
return self.fuzzy_symbol_hashed_equities.get(fuzzy_symbol, ())
return self._fuzzy_symbol_cache.get(fuzzy_symbol, ())
def _get_fuzzy_candidates_in_range(self, fuzzy_symbol, ad_value):
return only_active_assets(
@@ -975,7 +995,7 @@ class AssetFinderCachedEquities(AssetFinder):
)
def _get_split_candidates(self, company_symbol, share_class_symbol):
return self.company_share_class_hashed_equities.get(
return self._company_share_class_cache.get(
(company_symbol, share_class_symbol),
(),
)
@@ -998,7 +1018,8 @@ class AssetFinderCachedEquities(AssetFinder):
share_class_symbol,
ad_value):
equities = self._get_split_candidates(
company_symbol, share_class_symbol
company_symbol,
share_class_symbol
)
partial_candidates = []
for equity in equities:
+257
View File
@@ -0,0 +1,257 @@
from itertools import product
from string import ascii_uppercase
import pandas as pd
from pandas.tseries.offsets import MonthBegin
from six import iteritems
from .futures import CME_CODE_TO_MONTH
def make_rotating_equity_info(num_assets,
first_start,
frequency,
periods_between_starts,
asset_lifetime):
"""
Create a DataFrame representing lifetimes of assets that are constantly
rotating in and out of existence.
Parameters
----------
num_assets : int
How many assets to create.
first_start : pd.Timestamp
The start date for the first asset.
frequency : str or pd.tseries.offsets.Offset (e.g. trading_day)
Frequency used to interpret next two arguments.
periods_between_starts : int
Create a new asset every `frequency` * `periods_between_new`
asset_lifetime : int
Each asset exists for `frequency` * `asset_lifetime` days.
Returns
-------
info : pd.DataFrame
DataFrame representing newly-created assets.
"""
return pd.DataFrame(
{
'symbol': [chr(ord('A') + i) for i in range(num_assets)],
# Start a new asset every `periods_between_starts` days.
'start_date': pd.date_range(
first_start,
freq=(periods_between_starts * frequency),
periods=num_assets,
),
# Each asset lasts for `asset_lifetime` days.
'end_date': pd.date_range(
first_start + (asset_lifetime * frequency),
freq=(periods_between_starts * frequency),
periods=num_assets,
),
'exchange': 'TEST',
},
index=range(num_assets),
)
def make_simple_equity_info(sids,
start_date,
end_date,
symbols=None):
"""
Create a DataFrame representing assets that exist for the full duration
between `start_date` and `end_date`.
Parameters
----------
sids : array-like of int
start_date : pd.Timestamp, optional
end_date : pd.Timestamp, optional
symbols : list, optional
Symbols to use for the assets.
If not provided, symbols are generated from the sequence 'A', 'B', ...
Returns
-------
info : pd.DataFrame
DataFrame representing newly-created assets.
"""
num_assets = len(sids)
if symbols is None:
symbols = list(ascii_uppercase[:num_assets])
return pd.DataFrame(
{
'symbol': list(symbols),
'start_date': pd.to_datetime([start_date] * num_assets),
'end_date': pd.to_datetime([end_date] * num_assets),
'exchange': 'TEST',
},
index=sids,
columns=(
'start_date',
'end_date',
'symbol',
'exchange',
),
)
def make_jagged_equity_info(num_assets,
start_date,
first_end,
frequency,
periods_between_ends,
auto_close_delta):
"""
Create a DataFrame representing assets that all begin at the same start
date, but have cascading end dates.
Parameters
----------
num_assets : int
How many assets to create.
start_date : pd.Timestamp
The start date for all the assets.
first_end : pd.Timestamp
The date at which the first equity will end.
frequency : str or pd.tseries.offsets.Offset (e.g. trading_day)
Frequency used to interpret the next argument.
periods_between_ends : int
Starting after the first end date, end each asset every
`frequency` * `periods_between_ends`.
Returns
-------
info : pd.DataFrame
DataFrame representing newly-created assets.
"""
frame = pd.DataFrame(
{
'symbol': [chr(ord('A') + i) for i in range(num_assets)],
'start_date': start_date,
'end_date': pd.date_range(
first_end,
freq=(periods_between_ends * frequency),
periods=num_assets,
),
'exchange': 'TEST',
},
index=range(num_assets),
)
# Explicitly pass None to disable setting the auto_close_date column.
if auto_close_delta is not None:
frame['auto_close_date'] = frame['end_date'] + auto_close_delta
return frame
def make_future_info(first_sid,
root_symbols,
years,
notice_date_func,
expiration_date_func,
start_date_func,
month_codes=None):
"""
Create a DataFrame representing futures for `root_symbols` during `year`.
Generates a contract per triple of (symbol, year, month) supplied to
`root_symbols`, `years`, and `month_codes`.
Parameters
----------
first_sid : int
The first sid to use for assigning sids to the created contracts.
root_symbols : list[str]
A list of root symbols for which to create futures.
years : list[int or str]
Years (e.g. 2014), for which to produce individual contracts.
notice_date_func : (Timestamp) -> Timestamp
Function to generate notice dates from first of the month associated
with asset month code. Return NaT to simulate futures with no notice
date.
expiration_date_func : (Timestamp) -> Timestamp
Function to generate expiration dates from first of the month
associated with asset month code.
start_date_func : (Timestamp) -> Timestamp, optional
Function to generate start dates from first of the month associated
with each asset month code. Defaults to a start_date one year prior
to the month_code date.
month_codes : dict[str -> [1..12]], optional
Dictionary of month codes for which to create contracts. Entries
should be strings mapped to values from 1 (January) to 12 (December).
Default is zipline.futures.CME_CODE_TO_MONTH
Returns
-------
futures_info : pd.DataFrame
DataFrame of futures data suitable for passing to an AssetDBWriter.
"""
if month_codes is None:
month_codes = CME_CODE_TO_MONTH
year_strs = list(map(str, years))
years = [pd.Timestamp(s, tz='UTC') for s in year_strs]
# Pairs of string/date like ('K06', 2006-05-01)
contract_suffix_to_beginning_of_month = tuple(
(month_code + year_str[-2:], year + MonthBegin(month_num))
for ((year, year_str), (month_code, month_num))
in product(
zip(years, year_strs),
iteritems(month_codes),
)
)
contracts = []
parts = product(root_symbols, contract_suffix_to_beginning_of_month)
for sid, (root_sym, (suffix, month_begin)) in enumerate(parts, first_sid):
contracts.append({
'sid': sid,
'root_symbol': root_sym,
'symbol': root_sym + suffix,
'start_date': start_date_func(month_begin),
'notice_date': notice_date_func(month_begin),
'expiration_date': notice_date_func(month_begin),
'multiplier': 500,
})
return pd.DataFrame.from_records(contracts, index='sid').convert_objects()
def make_commodity_future_info(first_sid,
root_symbols,
years,
month_codes=None):
"""
Make futures testing data that simulates the notice/expiration date
behavior of physical commodities like oil.
Parameters
----------
first_sid : int
root_symbols : list[str]
years : list[int]
month_codes : dict[str -> int]
Expiration dates are on the 20th of the month prior to the month code.
Notice dates are are on the 20th two months prior to the month code.
Start dates are one year before the contract month.
See Also
--------
make_future_info
"""
nineteen_days = pd.Timedelta(days=19)
one_year = pd.Timedelta(days=365)
return make_future_info(
first_sid=first_sid,
root_symbols=root_symbols,
years=years,
notice_date_func=lambda dt: dt - MonthBegin(2) + nineteen_days,
expiration_date_func=lambda dt: dt - MonthBegin(1) + nineteen_days,
start_date_func=lambda dt: dt - one_year,
month_codes=month_codes,
)
+6
View File
@@ -162,6 +162,12 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
Dates for which adjustments are needed
assets : pd.Int64Index
Assets for which adjustments are needed.
Returns
-------
adjustments : list[dict[int -> Adjustment]]
A list of mappings from index to adjustment objects to apply at that
index.
"""
cdef int start_date = int((dates[0] - EPOCH).total_seconds())
+14 -13
View File
@@ -11,18 +11,16 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
from os.path import join
from textwrap import dedent
import bcolz
from bcolz import ctable
from intervaltree import IntervalTree
from numpy import nan_to_num
from os.path import join
import json
import os
import numpy as np
import pandas as pd
from zipline.gens.sim_engine import NANOS_IN_MINUTE
from zipline.data._minute_bar_internal import (
minute_value,
@@ -30,6 +28,7 @@ from zipline.data._minute_bar_internal import (
find_last_traded_position_internal
)
from zipline.gens.sim_engine import NANOS_IN_MINUTE
from zipline.utils.memoize import lazyval
US_EQUITIES_MINUTES_PER_DAY = 390
@@ -563,14 +562,16 @@ class BcolzMinuteBarWriter(object):
dts.astype('datetime64[ns]'))
ohlc_ratio = self._ohlc_ratio
open_col[dt_ixs] = (nan_to_num(cols['open']) * ohlc_ratio).\
astype(np.uint32)
high_col[dt_ixs] = (nan_to_num(cols['high']) * ohlc_ratio).\
astype(np.uint32)
low_col[dt_ixs] = (nan_to_num(cols['low']) * ohlc_ratio).\
astype(np.uint32)
close_col[dt_ixs] = (nan_to_num(cols['close']) * ohlc_ratio).\
astype(np.uint32)
def convert_col(col):
"""Adapt float column into a uint32 column.
"""
return (np.nan_to_num(col) * ohlc_ratio).astype(np.uint32)
open_col[dt_ixs] = convert_col(cols['open'])
high_col[dt_ixs] = convert_col(cols['high'])
low_col[dt_ixs] = convert_col(cols['low'])
close_col[dt_ixs] = convert_col(cols['close'])
vol_col[dt_ixs] = cols['volume'].astype(np.uint32)
table.append([
+344 -296
View File
@@ -11,15 +11,13 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import (
ABCMeta,
abstractmethod,
abstractproperty,
)
from abc import ABCMeta, abstractmethod, abstractproperty
from errno import ENOENT
from functools import partial
from os import remove
from os.path import exists
import sqlite3
import warnings
from bcolz import (
carray,
@@ -27,12 +25,12 @@ from bcolz import (
open as open_ctable,
)
from collections import namedtuple
from click import progressbar
import logbook
import numpy as np
from numpy import (
array,
int64,
float64,
floating,
full,
iinfo,
integer,
@@ -49,40 +47,39 @@ from pandas import (
NaT,
isnull,
)
from pandas.tslib import iNaT
from six import (
iteritems,
with_metaclass,
viewkeys,
)
from zipline.utils.input_validation import coerce_string, preprocess
from zipline.utils.functional import apply
from zipline.utils.input_validation import (
coerce_string,
preprocess,
expect_element,
)
from zipline.utils.sqlite_utils import group_into_chunks
from zipline.utils.memoize import lazyval
from zipline.utils.cli import maybe_show_progress
from ._equities import _compute_row_slices, _read_bcolz_data
from ._adjustments import load_adjustments_from_sqlite
import logbook
logger = logbook.Logger('UsEquityPricing')
OHLC = frozenset(['open', 'high', 'low', 'close'])
US_EQUITY_PRICING_BCOLZ_COLUMNS = [
US_EQUITY_PRICING_BCOLZ_COLUMNS = (
'open', 'high', 'low', 'close', 'volume', 'day', 'id'
]
SQLITE_ADJUSTMENT_COLUMNS = frozenset(['effective_date', 'ratio', 'sid'])
)
SQLITE_ADJUSTMENT_COLUMN_DTYPES = {
'effective_date': integer,
'ratio': floating,
'ratio': float,
'sid': integer,
}
SQLITE_ADJUSTMENT_TABLENAMES = frozenset(['splits', 'dividends', 'mergers'])
SQLITE_DIVIDEND_PAYOUT_COLUMNS = frozenset(
['sid',
'ex_date',
'declared_date',
'pay_date',
'record_date',
'amount'])
SQLITE_DIVIDEND_PAYOUT_COLUMN_DTYPES = {
'sid': integer,
'ex_date': integer,
@@ -92,15 +89,6 @@ SQLITE_DIVIDEND_PAYOUT_COLUMN_DTYPES = {
'amount': float,
}
SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMNS = frozenset(
['sid',
'ex_date',
'declared_date',
'record_date',
'pay_date',
'payment_sid',
'ratio'])
SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMN_DTYPES = {
'sid': integer,
'ex_date': integer,
@@ -120,75 +108,179 @@ class NoDataOnDate(Exception):
pass
class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
def check_uint32_safe(value, colname):
if value >= UINT32_MAX:
raise ValueError(
"Value %s from column '%s' is too large" % (value, colname)
)
@expect_element(invalid_data_behavior={'warn', 'raise', 'ignore'})
def winsorise_uint32(df, invalid_data_behavior, column, *columns):
"""Drops any record where a value would not fit into a uint32.
Parameters
----------
df : pd.DataFrame
The dataframe to winsorise.
invalid_data_behavior : {'warn', 'raise', 'ignore'}
What to do when data is outside the bounds of a uint32.
*columns : iterable[str]
The names of the columns to check.
Returns
-------
truncated : pd.DataFrame
``df`` with values that do not fit into a uint32 zeroed out.
"""
columns = list((column,) + columns)
mask = df[columns] > UINT32_MAX
if invalid_data_behavior != 'ignore':
mask |= df[columns].isnull()
else:
# we are not going to generate a warning or error for this so just use
# nan_to_num
df[columns] = np.nan_to_num(df[columns])
mv = mask.values
if mv.any():
if invalid_data_behavior == 'raise':
raise ValueError(
'%d values out of bounds for uint32: %r' % (
mv.sum(), df[mask.any(axis=1)],
),
)
if invalid_data_behavior == 'warn':
warnings.warn(
'Ignoring %d values because they are out of bounds for'
' uint32: %r' % (
mv.sum(), df[mask.any(axis=1)],
),
stacklevel=3, # one extra frame for `expect_element`
)
df[mask] = 0
return df
@expect_element(invalid_data_behavior={'warn', 'raise', 'ignore'})
def to_ctable(raw_data, invalid_data_behavior):
if isinstance(raw_data, ctable):
# we already have a ctable so do nothing
return raw_data
winsorise_uint32(raw_data, invalid_data_behavior, 'volume', *OHLC)
processed = (raw_data[list(OHLC)] * 1000).astype('uint32')
dates = raw_data.index.values.astype('datetime64[s]')
check_uint32_safe(dates.max().view(np.int64), 'day')
processed['day'] = dates.astype('uint32')
processed['volume'] = raw_data.volume.astype('uint32')
return ctable.fromdataframe(processed)
class BcolzDailyBarWriter(object):
"""
Class capable of writing daily OHLCV data to disk in a format that can be
read efficiently by BcolzDailyOHLCVReader.
Parameters
----------
filename : str
The location at which we should write our output.
calendar : pandas.DatetimeIndex
Calendar to use to compute asset calendar offsets.
See Also
--------
BcolzDailyBarReader : Consumer of the data written by this class.
"""
@abstractmethod
def gen_tables(self, assets):
"""
Return an iterator of pairs of (asset_id, bcolz.ctable).
"""
raise NotImplementedError()
_csv_dtypes = {
'open': float64,
'high': float64,
'low': float64,
'close': float64,
'volume': float64,
}
@abstractmethod
def to_uint32(self, array, colname):
"""
Convert raw column values produced by gen_tables into uint32 values.
def __init__(self, filename, calendar):
self._filename = filename
self._calendar = calendar
Parameters
----------
array : np.array
An array of raw values.
colname : str, {'open', 'high', 'low', 'close', 'volume', 'day'}
The name of the column being loaded.
@property
def progress_bar_message(self):
return "Merging asset files:"
For output being read by the default BcolzOHLCVReader, data should be
stored in the following manner:
def progress_bar_item_show_func(self, value):
return value if value is None else str(value[0])
- Pricing columns (Open, High, Low, Close) should be stored as 1000 *
as-traded dollar value.
- Volume should be the as-traded volume.
- Dates should be stored as seconds since midnight UTC, Jan 1, 1970.
"""
raise NotImplementedError()
def write(self, filename, calendar, assets, show_progress=False):
def write(self,
data,
assets=None,
show_progress=False,
invalid_data_behavior='warn'):
"""
Parameters
----------
filename : str
The location at which we should write our output.
calendar : pandas.DatetimeIndex
Calendar to use to compute asset calendar offsets.
assets : pandas.Int64Index
The assets for which to write data.
data : iterable[tuple[int, pandas.DataFrame or bcolz.ctable]]
The data chunks to write. Each chunk should be a tuple of sid
and the data for that asset.
assets : set[int], optional
The assets that should be in ``data``. If this is provided
we will check ``data`` against the assets and provide better
progress information.
show_progress : bool
Whether or not to show a progress bar while writing.
invalid_data_behavior : {'warn', 'raise', 'ignore'}
What to do when data is encountered that is outside the range of
a uint32.
Returns
-------
table : bcolz.ctable
The newly-written table.
"""
_iterator = self.gen_tables(assets)
if show_progress:
pbar = progressbar(
_iterator,
length=len(assets),
item_show_func=lambda i: i if i is None else str(i[0]),
label="Merging asset files:",
)
with pbar as pbar_iterator:
return self._write_internal(filename, calendar, pbar_iterator)
return self._write_internal(filename, calendar, _iterator)
ctx = maybe_show_progress(
((sid, to_ctable(df, invalid_data_behavior)) for sid, df in data),
show_progress=show_progress,
item_show_func=self.progress_bar_item_show_func,
label=self.progress_bar_message,
length=len(assets) if assets is not None else None,
)
with ctx as it:
return self._write_internal(it, assets)
def _write_internal(self, filename, calendar, iterator):
def write_csvs(self,
asset_map,
show_progress=False,
invalid_data_behavior='warn'):
"""Read CSVs as DataFrames from our asset map.
Parameters
----------
asset_map : dict[int -> str]
A mapping from asset id to file path with the CSV data for that
asset
show_progress : bool
Whether or not to show a progress bar while writing.
invalid_data_behavior : {'warn', 'raise', 'ignore'}
What to do when data is encountered that is outside the range of
a uint32.
"""
read = partial(
read_csv,
parse_dates=['day'],
index_col='day',
dtype=self._csv_dtypes,
)
return self.write(
((asset, read(path)) for asset, path in iteritems(asset_map)),
assets=viewkeys(asset_map),
show_progress=show_progress,
invalid_data_behavior=invalid_data_behavior,
)
def _write_internal(self, iterator, assets):
"""
Internal implementation of write.
@@ -206,6 +298,15 @@ class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
}
earliest_date = None
calendar = self._calendar
if assets is not None:
@apply
def iterator(iterator=iterator, assets=set(assets)):
for asset_id, table in iterator:
if asset_id not in assets:
raise ValueError('unknown asset id %r' % asset_id)
yield asset_id, table
for asset_id, table in iterator:
nrows = len(table)
@@ -213,11 +314,12 @@ class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
if column_name == 'id':
# We know what the content of this column is, so don't
# bother reading it.
columns['id'].append(full((nrows,), asset_id, uint32))
columns['id'].append(
full((nrows,), asset_id, dtype='uint32'),
)
continue
columns[column_name].append(
self.to_uint32(table[column_name][:], column_name)
)
columns[column_name].append(table[column_name])
if earliest_date is None:
earliest_date = table["day"][0]
@@ -238,11 +340,7 @@ class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
# Calculate the number of trading days between the first date
# in the stored data and the first date of **this** asset. This
# offset used for output alignment by the reader.
# HACK: Index with a list so that we get back an array we can pass
# to self.to_uint32. We could try to extract this in the loop
# above, but that makes the logic a lot messier.
asset_first_day = self.to_uint32(table['day'][[0]], 'day')[0]
asset_first_day = table['day'][0]
calendar_offset[asset_key] = calendar.get_loc(
Timestamp(asset_first_day, unit='s', tz='UTC'),
)
@@ -254,12 +352,15 @@ class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
for colname in US_EQUITY_PRICING_BCOLZ_COLUMNS
],
names=US_EQUITY_PRICING_BCOLZ_COLUMNS,
rootdir=filename,
rootdir=self._filename,
mode='w',
)
full_table.attrs['first_trading_day'] = \
int(earliest_date / 1e6)
full_table.attrs['first_trading_day'] = (
earliest_date // 1e6
if earliest_date is not None else
iNaT
)
full_table.attrs['first_row'] = first_row
full_table.attrs['last_row'] = last_row
full_table.attrs['calendar_offset'] = calendar_offset
@@ -267,68 +368,6 @@ class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
return full_table
class DailyBarWriterFromCSVs(BcolzDailyBarWriter):
"""
BcolzDailyBarWriter constructed from a map from csvs to assets.
Parameters
----------
asset_map : dict
A map from asset_id -> path to csv with data for that asset.
CSVs should have the following columns:
day : datetime64
open : float64
high : float64
low : float64
close : float64
volume : int64
"""
_csv_dtypes = {
'open': float64,
'high': float64,
'low': float64,
'close': float64,
'volume': float64,
}
def __init__(self, asset_map):
self._asset_map = asset_map
def gen_tables(self, assets):
"""
Read CSVs as DataFrames from our asset map.
"""
dtypes = self._csv_dtypes
for asset in assets:
path = self._asset_map.get(asset)
if path is None:
raise KeyError("No path supplied for asset %s" % asset)
data = read_csv(path, parse_dates=['day'], dtype=dtypes)
yield asset, ctable.fromdataframe(data)
def to_uint32(self, array, colname):
arrmax = array.max()
if colname in OHLC:
self.check_uint_safe(arrmax * 1000, colname)
return (array * 1000).astype(uint32)
elif colname == 'volume':
self.check_uint_safe(arrmax, colname)
return array.astype(uint32)
elif colname == 'day':
nanos_per_second = (1000 * 1000 * 1000)
self.check_uint_safe(arrmax.view(int64) / nanos_per_second,
colname)
return (array.view(int64) / nanos_per_second).astype(uint32)
@staticmethod
def check_uint_safe(value, colname):
if value >= UINT32_MAX:
raise ValueError(
"Value %s from column '%s' is too large" % (value, colname)
)
class DailyBarReader(with_metaclass(ABCMeta)):
"""
Reader for OHCLV pricing data at a daily frequency.
@@ -402,37 +441,59 @@ class BcolzDailyBarReader(DailyBarReader):
def __init__(self, table):
self._table = table
self._calendar = DatetimeIndex(table.attrs['calendar'], tz='UTC')
self._first_rows = {
int(asset_id): start_index
for asset_id, start_index in iteritems(table.attrs['first_row'])
}
self._last_rows = {
int(asset_id): end_index
for asset_id, end_index in iteritems(table.attrs['last_row'])
}
self._calendar_offsets = {
int(id_): offset
for id_, offset in iteritems(table.attrs['calendar_offset'])
}
try:
self._first_trading_day = Timestamp(
table.attrs['first_trading_day'],
unit='ms',
tz='UTC'
)
except KeyError:
self._first_trading_day = None
# Cache of fully read np.array for the carrays in the daily bar table.
# raw_array does not use the same cache, but it could.
# Need to test keeping the entire array in memory for the course of a
# process first.
self._spot_cols = {}
self.PRICE_ADJUSTMENT_FACTOR = 0.001
@lazyval
def _calendar(self):
return DatetimeIndex(self._table.attrs['calendar'], tz='UTC')
@lazyval
def _first_rows(self):
return {
int(asset_id): start_index
for asset_id, start_index in iteritems(
self._table.attrs['first_row'],
)
}
@lazyval
def _last_rows(self):
return {
int(asset_id): end_index
for asset_id, end_index in iteritems(
self._table.attrs['last_row'],
)
}
@lazyval
def _calendar_offsets(self):
return {
int(id_): offset
for id_, offset in iteritems(
self._table.attrs['calendar_offset'],
)
}
@lazyval
def first_trading_day(self):
try:
return Timestamp(
self._table.attrs['first_trading_day'],
unit='ms',
tz='UTC'
)
except KeyError:
return None
@property
def last_available_dt(self):
return self._calendar[-1]
def _compute_slices(self, start_idx, end_idx, assets):
"""
Compute the raw row indices to load for each asset on a query for the
@@ -493,14 +554,6 @@ class BcolzDailyBarReader(DailyBarReader):
offsets,
)
@property
def first_trading_day(self):
return self._first_trading_day
@property
def last_available_dt(self):
return self._calendar[-1]
def _spot_col(self, colname):
"""
Get the colname from daily_bar_table and read all of it into memory,
@@ -713,6 +766,8 @@ class SQLiteAdjustmentWriter(object):
----------
conn_or_path : str or sqlite3.Connection
A handle to the target sqlite database.
daily_bar_reader : BcolzDailyBarReader
Daily bar reader to use for dividend writes.
overwrite : bool, optional, default=False
If True and conn_or_path is a string, remove any existing files at the
given path before connecting.
@@ -722,7 +777,10 @@ class SQLiteAdjustmentWriter(object):
SQLiteAdjustmentReader
"""
def __init__(self, conn_or_path, calendar, daily_bar_reader,
def __init__(self,
conn_or_path,
daily_bar_reader,
calendar,
overwrite=False):
if isinstance(conn_or_path, sqlite3.Connection):
self.conn = conn_or_path
@@ -734,98 +792,80 @@ class SQLiteAdjustmentWriter(object):
if e.errno != ENOENT:
raise
self.conn = sqlite3.connect(conn_or_path)
self.uri = conn_or_path
else:
raise TypeError("Unknown connection type %s" % type(conn_or_path))
self._daily_bar_reader = daily_bar_reader
self._calendar = calendar
def write_frame(self, tablename, frame):
if frozenset(frame.columns) != SQLITE_ADJUSTMENT_COLUMNS:
raise ValueError(
"Unexpected frame columns:\n"
"Expected Columns: %s\n"
"Received Columns: %s" % (
SQLITE_ADJUSTMENT_COLUMNS,
frame.columns.tolist(),
)
def _write(self, tablename, expected_dtypes, frame):
if frame is None or frame.empty:
# keeping the dtypes correct for empty frames is not easy
frame = DataFrame(
np.array([], dtype=list(expected_dtypes.items())),
)
elif tablename not in SQLITE_ADJUSTMENT_TABLENAMES:
raise ValueError(
"Adjustment table %s not in %s" % (
tablename, SQLITE_ADJUSTMENT_TABLENAMES
)
)
expected_dtypes = SQLITE_ADJUSTMENT_COLUMN_DTYPES
actual_dtypes = frame.dtypes
for colname, expected in iteritems(expected_dtypes):
actual = actual_dtypes[colname]
if not issubdtype(actual, expected):
raise TypeError(
"Expected data of type {expected} for column '{colname}', "
"but got {actual}.".format(
expected=expected,
colname=colname,
actual=actual,
else:
if frozenset(frame.columns) != viewkeys(expected_dtypes):
raise ValueError(
"Unexpected frame columns:\n"
"Expected Columns: %s\n"
"Received Columns: %s" % (
set(expected_dtypes),
frame.columns.tolist(),
)
)
return frame.to_sql(tablename, self.conn)
actual_dtypes = frame.dtypes
for colname, expected in iteritems(expected_dtypes):
actual = actual_dtypes[colname]
if not issubdtype(actual, expected):
raise TypeError(
"Expected data of type {expected} for column"
" '{colname}', but got '{actual}'.".format(
expected=expected,
colname=colname,
actual=actual,
),
)
frame.to_sql(
tablename,
self.conn,
if_exists='append',
chunksize=50000,
)
def write_frame(self, tablename, frame):
if tablename not in SQLITE_ADJUSTMENT_TABLENAMES:
raise ValueError(
"Adjustment table %s not in %s" % (
tablename,
SQLITE_ADJUSTMENT_TABLENAMES,
)
)
return self._write(
tablename,
SQLITE_ADJUSTMENT_COLUMN_DTYPES,
frame,
)
def write_dividend_payouts(self, frame):
"""
Write dividend payout data to SQLite table `dividend_payouts`.
"""
if frozenset(frame.columns) != SQLITE_DIVIDEND_PAYOUT_COLUMNS:
raise ValueError(
"Unexpected frame columns:\n"
"Expected Columns: %s\n"
"Received Columns: %s" % (
sorted(SQLITE_DIVIDEND_PAYOUT_COLUMNS),
sorted(frame.columns.tolist()),
)
)
expected_dtypes = SQLITE_DIVIDEND_PAYOUT_COLUMN_DTYPES
actual_dtypes = frame.dtypes
for colname, expected in iteritems(expected_dtypes):
actual = actual_dtypes[colname]
if not issubdtype(actual, expected):
raise TypeError(
"Expected data of type {expected} for column '{colname}', "
"but got {actual}.".format(
expected=expected,
colname=colname,
actual=actual,
)
)
return frame.to_sql('dividend_payouts', self.conn)
return self._write(
'dividend_payouts',
SQLITE_DIVIDEND_PAYOUT_COLUMN_DTYPES,
frame,
)
def write_stock_dividend_payouts(self, frame):
if frozenset(frame.columns) != SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMNS:
raise ValueError(
"Unexpected frame columns:\n"
"Expected Columns: %s\n"
"Received Columns: %s" % (
sorted(SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMNS),
sorted(frame.columns.tolist()),
)
)
expected_dtypes = SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMN_DTYPES
actual_dtypes = frame.dtypes
for colname, expected in iteritems(expected_dtypes):
actual = actual_dtypes[colname]
if not issubdtype(actual, expected):
raise TypeError(
"Expected data of type {expected} for column '{colname}', "
"but got {actual}.".format(
expected=expected,
colname=colname,
actual=actual,
)
)
return frame.to_sql('stock_dividend_payouts', self.conn)
return self._write(
'stock_dividend_payouts',
SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMN_DTYPES,
frame,
)
def calc_dividend_ratios(self, dividends):
"""
@@ -841,6 +881,15 @@ class SQLiteAdjustmentWriter(object):
- effective_date, the date in seconds on which to apply the ratio.
- ratio, the ratio to apply to backwards looking pricing data.
"""
if dividends is None:
return DataFrame(np.array(
[],
dtype=[
('sid', uint32),
('effective_date', uint32),
('ratio', float64),
],
))
ex_dates = dividends.ex_date.values
sids = dividends.sid.values
@@ -850,14 +899,12 @@ class SQLiteAdjustmentWriter(object):
daily_bar_reader = self._daily_bar_reader
calendar = self._calendar
effective_dates = full(len(amounts), -1, dtype=int64)
calendar = self._calendar
for i, amount in enumerate(amounts):
sid = sids[i]
ex_date = ex_dates[i]
day_loc = calendar.get_loc(ex_date)
day_loc = calendar.get_loc(ex_date, method='bfill')
prev_close_date = calendar[day_loc - 1]
try:
prev_close = daily_bar_reader.spot_price(
@@ -890,28 +937,29 @@ class SQLiteAdjustmentWriter(object):
'ratio': ratios,
})
def write_dividend_data(self, dividends, stock_dividends=None):
"""
Write both dividend payouts and the derived price adjustment ratios.
"""
# First write the dividend payouts.
dividend_payouts = dividends.copy()
dividend_payouts['ex_date'] = dividend_payouts['ex_date'].values.\
astype('datetime64[s]').astype(integer)
dividend_payouts['record_date'] = \
dividend_payouts['record_date'].values.astype('datetime64[s]').\
astype(integer)
dividend_payouts['declared_date'] = \
dividend_payouts['declared_date'].values.astype('datetime64[s]').\
astype(integer)
dividend_payouts['pay_date'] = \
dividend_payouts['pay_date'].values.astype('datetime64[s]').\
astype(integer)
def _write_dividends(self, dividends):
if dividends is None:
dividend_payouts = None
else:
dividend_payouts = dividends.copy()
dividend_payouts['ex_date'] = dividend_payouts['ex_date'].values.\
astype('datetime64[s]').astype(integer)
dividend_payouts['record_date'] = \
dividend_payouts['record_date'].values.astype('datetime64[s]').\
astype(integer)
dividend_payouts['declared_date'] = \
dividend_payouts['declared_date'].values.astype('datetime64[s]').\
astype(integer)
dividend_payouts['pay_date'] = \
dividend_payouts['pay_date'].values.astype('datetime64[s]').\
astype(integer)
self.write_dividend_payouts(dividend_payouts)
if stock_dividends is not None:
def _write_stock_dividends(self, stock_dividends):
if stock_dividends is None:
stock_dividend_payouts = None
else:
stock_dividend_payouts = stock_dividends.copy()
stock_dividend_payouts['ex_date'] = \
stock_dividend_payouts['ex_date'].values.\
@@ -925,26 +973,26 @@ class SQLiteAdjustmentWriter(object):
stock_dividend_payouts['pay_date'] = \
stock_dividend_payouts['pay_date'].\
values.astype('datetime64[s]').astype(integer)
else:
stock_dividend_payouts = DataFrame({
'sid': array([], dtype=uint32),
'record_date': array([], dtype=uint32),
'ex_date': array([], dtype=uint32),
'declared_date': array([], dtype=uint32),
'pay_date': array([], dtype=uint32),
'payment_sid': array([], dtype=uint32),
'ratio': array([], dtype=float),
})
self.write_stock_dividend_payouts(stock_dividend_payouts)
def write_dividend_data(self, dividends, stock_dividends=None):
"""
Write both dividend payouts and the derived price adjustment ratios.
"""
# First write the dividend payouts.
self._write_dividends(dividends)
self._write_stock_dividends(stock_dividends)
# Second from the dividend payouts, calculate ratios.
dividend_ratios = self.calc_dividend_ratios(dividends)
self.write_frame('dividends', dividend_ratios)
def write(self, splits, mergers, dividends, stock_dividends=None):
def write(self,
splits=None,
mergers=None,
dividends=None,
stock_dividends=None):
"""
Writes data to a SQLite file to be read by SQLiteAdjustmentReader.
+21
View File
@@ -0,0 +1,21 @@
"""dispatcher object with a custom namespace.
Anything that has been dispatched will also be put into this module.
"""
from functools import partial
import sys
from multipledispatch import dispatch
try:
from datashape.dispatch import namespace
except ImportError:
pass
else:
globals().update(namespace)
dispatch = partial(dispatch, namespace=globals())
del namespace
del partial
del sys
+1 -2
View File
@@ -46,6 +46,5 @@ if __name__ == '__main__':
end=end,
)
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data,
identifiers=stocks)
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data)
results = algo.run(input_data)
+1 -2
View File
@@ -60,8 +60,7 @@ if __name__ == '__main__':
end=end)
# Create and run the algorithm.
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data,
identifiers=['AAPL'])
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data)
results = algo.run(data)
analyze(results=results)
+1 -2
View File
@@ -115,8 +115,7 @@ if __name__ == '__main__':
end=end)
# Create and run the algorithm.
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data,
identifiers=['AAPL'])
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data)
results = algo.run(data).dropna()
# Plot the portfolio and asset data.
+1 -2
View File
@@ -115,8 +115,7 @@ if __name__ == '__main__':
end=end)
# Create and run the algorithm.
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data,
identifiers=['AAPL'])
algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data)
results = algo.run(data)
# Plot the portfolio and asset data.
+1 -3
View File
@@ -173,9 +173,7 @@ if __name__ == '__main__':
data = data.dropna()
# Create and run the algorithm.
olmar = TradingAlgorithm(handle_data=handle_data,
initialize=initialize,
identifiers=STOCKS)
olmar = TradingAlgorithm(handle_data=handle_data, initialize=initialize)
results = olmar.run(data)
# Plot the portfolio data.
+56 -138
View File
@@ -22,13 +22,9 @@ import numpy as np
from six import string_types
from sqlalchemy import create_engine
from zipline.assets import AssetDBWriter, AssetFinder
from zipline.data.loader import load_market_data
from zipline.utils import tradingcalendar
from zipline.assets import AssetFinder
from zipline.assets.asset_writer import (
AssetDBWriterFromList,
AssetDBWriterFromDictionary,
AssetDBWriterFromDataFrame)
from zipline.errors import (
NoFurtherDataError
)
@@ -37,47 +33,61 @@ from zipline.utils.memoize import remember_last, lazyval
log = logbook.Logger('Trading')
# The financial simulations in zipline depend on information
# about the benchmark index and the risk free rates of return.
# The benchmark index defines the benchmark returns used in
# the calculation of performance metrics such as alpha/beta. Many
# components, including risk, performance, transforms, and
# batch_transforms, need access to a calendar of trading days and
# market hours. The TradingEnvironment maintains two time keeping
# facilities:
# - a DatetimeIndex of trading days for calendar calculations
# - a timezone name, which should be local to the exchange
# hosting the benchmark index. All dates are normalized to UTC
# for serialization and storage, and the timezone is used to
# ensure proper rollover through daylight savings and so on.
#
# User code will not normally need to use TradingEnvironment
# directly. If you are extending zipline's core financial
# components and need to use the environment, you must import the module and
# build a new TradingEnvironment object, then pass that TradingEnvironment as
# the 'env' arg to your TradingAlgorithm.
class TradingEnvironment(object):
"""
The financial simulations in zipline depend on information
about the benchmark index and the risk free rates of return.
The benchmark index defines the benchmark returns used in
the calculation of performance metrics such as alpha/beta. Many
components, including risk, performance, transforms, and
batch_transforms, need access to a calendar of trading days and
market hours. The TradingEnvironment maintains two time keeping
facilities:
- a DatetimeIndex of trading days for calendar calculations
- a timezone name, which should be local to the exchange
hosting the benchmark index. All dates are normalized to UTC
for serialization and storage, and the timezone is used to
ensure proper rollover through daylight savings and so on.
User code will not normally need to use TradingEnvironment
directly. If you are extending zipline's core financial
components and need to use the environment, you must import the module and
build a new TradingEnvironment object, then pass that TradingEnvironment as
the 'env' arg to your TradingAlgorithm.
Parameters
----------
load : callable, optional
The function that returns benchmark returns and treasury curves.
The treasury curves are expected to be a DataFrame with an index of
dates and columns of the curve names, e.g. '10year', '1month', etc.
bm_symbol : str, optional
The benchmark symbol
exchange_tz : tz-coercable, optional
The timezone of the exchange.
min_date : datetime, optional
The oldest date that we know about in this environment.
max_date : datetime, optional
The most recent date that we know about in this environment.
env_trading_calendar : pd.DatetimeIndex, optional
The calendar of datetimes that define our market hours.
asset_db_path : str or sa.engine.Engine, optional
The path to the assets db or sqlalchemy Engine object to use to
construct an AssetFinder.
"""
# Token used as a substitute for pickling objects that contain a
# reference to a TradingEnvironment
PERSISTENT_TOKEN = "<TradingEnvironment>"
def __init__(
self,
load=None,
bm_symbol='^GSPC',
exchange_tz="US/Eastern",
min_date=None,
max_date=None,
env_trading_calendar=tradingcalendar,
asset_db_path=':memory:'
):
"""
@load is function that returns benchmark_returns and treasury_curves
The treasury_curves are expected to be a DataFrame with an index of
dates and columns of the curve names, e.g. '10year', '1month', etc.
"""
def __init__(self,
load=None,
bm_symbol='^GSPC',
exchange_tz="US/Eastern",
min_date=None,
max_date=None,
env_trading_calendar=tradingcalendar,
asset_db_path=':memory:'):
self.trading_day = env_trading_calendar.trading_day.copy()
# `tc_td` is short for "trading calendar trading days"
@@ -114,11 +124,11 @@ class TradingEnvironment(object):
if isinstance(asset_db_path, string_types):
asset_db_path = 'sqlite:///%s' % asset_db_path
self.engine = engine = create_engine(asset_db_path)
AssetDBWriterFromDictionary().init_db(engine)
else:
self.engine = engine = asset_db_path
if engine is not None:
AssetDBWriter(engine).init_db()
self.asset_finder = AssetFinder(engine)
else:
self.asset_finder = None
@@ -128,107 +138,15 @@ class TradingEnvironment(object):
return self.minutes_for_days_in_range(self.first_trading_day,
self.last_trading_day)
def write_data(self,
engine=None,
equities_data=None,
futures_data=None,
exchanges_data=None,
root_symbols_data=None,
equities_df=None,
futures_df=None,
exchanges_df=None,
root_symbols_df=None,
equities_identifiers=None,
futures_identifiers=None,
exchanges_identifiers=None,
root_symbols_identifiers=None,
allow_sid_assignment=True):
""" Write the supplied data to the database.
def write_data(self, **kwargs):
"""Write data into the asset_db.
Parameters
----------
equities_data: dict, optional
A dictionary of equity metadata
futures_data: dict, optional
A dictionary of futures metadata
exchanges_data: dict, optional
A dictionary of exchanges metadata
root_symbols_data: dict, optional
A dictionary of root symbols metadata
equities_df: pandas.DataFrame, optional
A pandas.DataFrame of equity metadata
futures_df: pandas.DataFrame, optional
A pandas.DataFrame of futures metadata
exchanges_df: pandas.DataFrame, optional
A pandas.DataFrame of exchanges metadata
root_symbols_df: pandas.DataFrame, optional
A pandas.DataFrame of root symbols metadata
equities_identifiers: list, optional
A list of equities identifiers (sids, symbols, Assets)
futures_identifiers: list, optional
A list of futures identifiers (sids, symbols, Assets)
exchanges_identifiers: list, optional
A list of exchanges identifiers (ids or names)
root_symbols_identifiers: list, optional
A list of root symbols identifiers (ids or symbols)
**kwargs
Forwarded to AssetDBWriter.write
"""
if engine:
self.engine = engine
# If any pandas.DataFrame data has been provided,
# write it to the database.
has_rows = lambda df: df is not None and len(df) > 0
if any(map(has_rows, [equities_df,
futures_df,
exchanges_df,
root_symbols_df])):
self._write_data_dataframes(
equities=equities_df,
futures=futures_df,
exchanges=exchanges_df,
root_symbols=root_symbols_df,
)
# Same for dicts.
has_data = lambda d: d is not None and len(d) > 0
if any(map(has_data, [equities_data,
futures_data,
exchanges_data,
futures_data])):
self._write_data_dicts(
equities=equities_data,
futures=futures_data,
exchanges=exchanges_data,
root_symbols=root_symbols_data
)
# Same for iterables.
if any(map(has_data, [equities_identifiers,
futures_identifiers,
exchanges_identifiers,
root_symbols_identifiers])):
self._write_data_lists(
equities=equities_identifiers,
futures=futures_identifiers,
exchanges=exchanges_identifiers,
root_symbols=root_symbols_identifiers,
allow_sid_assignment=allow_sid_assignment
)
def _write_data_lists(self, equities=None, futures=None, exchanges=None,
root_symbols=None, allow_sid_assignment=True):
AssetDBWriterFromList(equities, futures, exchanges, root_symbols)\
.write_all(self.engine, allow_sid_assignment=allow_sid_assignment)
def _write_data_dicts(self, equities=None, futures=None, exchanges=None,
root_symbols=None):
AssetDBWriterFromDictionary(equities, futures, exchanges, root_symbols)\
.write_all(self.engine)
def _write_data_dataframes(self, equities=None, futures=None,
exchanges=None, root_symbols=None):
AssetDBWriterFromDataFrame(equities, futures, exchanges, root_symbols)\
.write_all(self.engine)
AssetDBWriter(self.engine).write(**kwargs)
def normalize_date(self, test_date):
test_date = pd.Timestamp(test_date, tz='UTC')
+86 -92
View File
@@ -1,8 +1,6 @@
"""
Synthetic data loaders for testing.
"""
from bcolz import ctable
from numpy import (
arange,
array,
@@ -20,7 +18,6 @@ from sqlite3 import connect as sqlite3_connect
from .base import PipelineLoader
from .frame import DataFrameLoader
from zipline.data.us_equity_pricing import (
BcolzDailyBarWriter,
SQLiteAdjustmentReader,
SQLiteAdjustmentWriter,
US_EQUITY_PRICING_BCOLZ_COLUMNS,
@@ -195,9 +192,29 @@ class SeededRandomLoader(PrecomputedLoader):
return self.state.randn(*shape) < 0
class SyntheticDailyBarWriter(BcolzDailyBarWriter):
OHLCV = ('open', 'high', 'low', 'close', 'volume')
OHLC = ('open', 'high', 'low', 'close')
PSEUDO_EPOCH = Timestamp('2000-01-01', tz='UTC')
def asset_start(asset_info, asset):
ret = asset_info.loc[asset]['start_date']
if ret.tz is None:
ret = ret.tz_localize('UTC')
assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp"
return ret
def asset_end(asset_info, asset):
ret = asset_info.loc[asset]['end_date']
if ret.tz is None:
ret = ret.tz_localize('UTC')
assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp"
return ret
def make_daily_bar_data(asset_info, calendar):
"""
Bcolz writer that creates synthetic data based on asset lifetime metadata.
For a given asset/date/column combination, we generate a corresponding raw
value using the following formula for OHLCV columns:
@@ -221,50 +238,50 @@ class SyntheticDailyBarWriter(BcolzDailyBarWriter):
----------
asset_info : DataFrame
DataFrame with asset_id as index and 'start_date'/'end_date' columns.
calendar : DatetimeIndex
Calendar to use for constructing asset lifetimes.
calendar : pd.DatetimeIndex
The trading calendar to use.
Yields
------
p : (int, pd.DataFrame)
A sid, data pair to be passed to BcolzDailyDailyBarWriter.write
"""
OHLCV = ('open', 'high', 'low', 'close', 'volume')
OHLC = ('open', 'high', 'low', 'close')
PSEUDO_EPOCH = Timestamp('2000-01-01', tz='UTC')
assert (
# Using .value here to avoid having to care about UTC-aware dates.
PSEUDO_EPOCH.value <
calendar.min().value <=
asset_info['start_date'].min().value
), "calendar.min(): %s\nasset_info['start_date'].min(): %s" % (
calendar.min(),
asset_info['start_date'].min(),
)
def __init__(self, asset_info, calendar):
super(SyntheticDailyBarWriter, self).__init__()
assert (
# Using .value here to avoid having to care about UTC-aware dates.
self.PSEUDO_EPOCH.value <
calendar.min().value <=
asset_info['start_date'].min().value
)
assert (asset_info['start_date'] < asset_info['end_date']).all()
self._asset_info = asset_info
self._calendar = calendar
assert (asset_info['start_date'] < asset_info['end_date']).all()
def _raw_data_for_asset(self, asset_id):
def _raw_data_for_asset(asset_id):
"""
Generate 'raw' data that encodes information about the asset.
See class docstring for a description of the data format.
See docstring for a description of the data format.
"""
# Get the dates for which this asset existed according to our asset
# info.
dates = self._calendar[
self._calendar.slice_indexer(
self.asset_start(asset_id), self.asset_end(asset_id)
)
]
dates = calendar[calendar.slice_indexer(
asset_start(asset_info, asset_id),
asset_end(asset_info, asset_id),
)]
data = full(
(len(dates), len(US_EQUITY_PRICING_BCOLZ_COLUMNS)),
asset_id * (100 * 1000),
asset_id * 100 * 1000,
dtype=uint32,
)
# Add 10,000 * column-index to OHLCV columns
data[:, :5] += arange(5, dtype=uint32) * (10 * 1000)
data[:, :5] += arange(5, dtype=uint32) * 1000
# Add days since Jan 1 2001 for OHLCV columns.
data[:, :5] += (dates - self.PSEUDO_EPOCH).days[:, None].astype(uint32)
data[:, :5] += (dates - PSEUDO_EPOCH).days[:, None].astype(uint32)
frame = DataFrame(
data,
@@ -274,76 +291,53 @@ class SyntheticDailyBarWriter(BcolzDailyBarWriter):
frame['day'] = nanos_to_seconds(dates.asi8)
frame['id'] = asset_id
return frame
return ctable.fromdataframe(frame)
for asset in asset_info.index:
yield asset, _raw_data_for_asset(asset)
def asset_start(self, asset):
ret = self._asset_info.loc[asset]['start_date']
if ret.tz is None:
ret = ret.tz_localize('UTC')
assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp"
return ret
def asset_end(self, asset):
ret = self._asset_info.loc[asset]['end_date']
if ret.tz is None:
ret = ret.tz_localize('UTC')
assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp"
return ret
def expected_daily_bar_value(asset_id, date, colname):
"""
Check that the raw value for an asset/date/column triple is as
expected.
@classmethod
def expected_value(cls, asset_id, date, colname):
"""
Check that the raw value for an asset/date/column triple is as
expected.
Used by tests to verify data written by a writer.
"""
from_asset = asset_id * 100000
from_colname = OHLCV.index(colname) * 1000
from_date = (date - PSEUDO_EPOCH).days
return from_asset + from_colname + from_date
Used by tests to verify data written by a writer.
"""
from_asset = asset_id * 100 * 1000
from_colname = cls.OHLCV.index(colname) * (10 * 1000)
from_date = (date - cls.PSEUDO_EPOCH).days
return from_asset + from_colname + from_date
def expected_values_2d(self, dates, assets, colname):
"""
Return an 2D array containing cls.expected_value(asset_id, date,
colname) for each date/asset pair in the inputs.
def expected_daily_bar_values_2d(dates, asset_info, colname):
"""
Return an 2D array containing cls.expected_value(asset_id, date,
colname) for each date/asset pair in the inputs.
Values before/after an assets lifetime are filled with 0 for volume and
NaN for price columns.
"""
if colname == 'volume':
dtype = uint32
missing = 0
else:
dtype = float64
missing = float('nan')
Values before/after an assets lifetime are filled with 0 for volume and
NaN for price columns.
"""
if colname == 'volume':
dtype = uint32
missing = 0
else:
dtype = float64
missing = float('nan')
data = full((len(dates), len(assets)), missing, dtype=dtype)
for j, asset in enumerate(assets):
start, end = self.asset_start(asset), self.asset_end(asset)
for i, date in enumerate(dates):
# No value expected for dates outside the asset's start/end
# date.
if not (start <= date <= end):
continue
data[i, j] = self.expected_value(asset, date, colname)
return data
assets = asset_info.index
# BEGIN SUPERCLASS INTERFACE
def gen_tables(self, assets):
for asset in assets:
yield asset, self._raw_data_for_asset(asset)
def to_uint32(self, array, colname):
if colname in {'open', 'high', 'low', 'close'}:
# Data is stored as 1000 * raw value.
assert array.max() < (UINT_32_MAX / 1000), "Test data overflow!"
return array * 1000
else:
assert colname in ('volume', 'day'), "Unknown column: %s" % colname
return array
# END SUPERCLASS INTERFACE
data = full((len(dates), len(assets)), missing, dtype=dtype)
for j, asset in enumerate(assets):
start = asset_start(asset_info, asset)
end = asset_end(asset_info, asset)
for i, date in enumerate(dates):
# No value expected for dates outside the asset's start/end
# date.
if not (start <= date <= end):
continue
data[i, j] = expected_daily_bar_value(asset, date, colname)
return data
class NullAdjustmentReader(SQLiteAdjustmentReader):
+5 -2
View File
@@ -1,7 +1,10 @@
from zipline.sources.data_frame_source import DataFrameSource, DataPanelSource
from zipline.sources.test_source import SpecificEquityTrades
from .data_source import DataSource
from .data_frame_source import DataFrameSource, DataPanelSource
from .test_source import SpecificEquityTrades
from .simulated import RandomWalkSource
__all__ = [
'DataSource',
'DataFrameSource',
'DataPanelSource',
'SpecificEquityTrades',
+23 -7
View File
@@ -2,6 +2,9 @@ from .core import ( # noqa
EPOCH,
ExceptionSource,
ExplodingObject,
FakeDataPortal,
FetcherDataPortal,
MockDailyBarReader,
add_security_data,
all_pairs_matching_predicate,
all_subindices,
@@ -10,27 +13,40 @@ from .core import ( # noqa
check_allclose,
check_arrays,
chrange,
create_daily_df_for_asset,
create_data_portal,
create_data_portal_from_trade_history,
create_empty_splits_mergers_frame,
create_minute_bar_data,
create_minute_df_for_asset,
drain_zipline,
empty_asset_finder,
empty_assets_db,
empty_trading_env,
gen_calendars,
make_commodity_future_info,
make_future_info,
make_jagged_equity_info,
make_rotating_equity_info,
make_simple_equity_info,
make_test_handler,
make_trade_data_for_asset_info,
parameter_space,
parameter_space,
patch_os_environment,
permute_rows,
powerset,
product_upper_triangle,
read_compressed,
seconds_to_timestamp,
security_list_copy,
setup_logger,
str_to_seconds,
subtest,
teardown_logger,
temp_pipeline_engine,
test_resource_path,
tmp_asset_finder,
tmp_assets_db,
tmp_bcolz_minute_bar_reader,
tmp_dir,
tmp_trading_env,
to_series,
to_utc,
trades_by_sid_to_dfs,
write_bcolz_minute_data,
write_compressed,
)
+395 -491
View File
File diff suppressed because it is too large Load Diff
+623 -117
View File
@@ -1,28 +1,64 @@
from abc import ABCMeta, abstractproperty
import gc
import sqlite3
from unittest import TestCase
from contextlib2 import ExitStack
from logbook import NullHandler
from logbook import NullHandler, Logger
from nose_parameterized import parameterized
import numpy as np
import pandas as pd
from six import iteritems
from testfixtures import TempDirectory
from pandas.util.testing import assert_series_equal
import responses
from six import with_metaclass, iteritems
from toolz import flip
from ..assets.synthetic import make_simple_equity_info
from .core import (
create_daily_bar_data,
create_minute_bar_data,
gen_calendars,
tmp_asset_finder,
tmp_dir,
)
from ..data.data_portal import DataPortal
from ..data.us_equity_pricing import (
SQLiteAdjustmentReader,
SQLiteAdjustmentWriter,
)
from ..finance.trading import TradingEnvironment
from ..data.us_equity_pricing import (
BcolzDailyBarReader,
BcolzDailyBarWriter,
)
from ..data.minute_bars import (
BcolzMinuteBarReader,
BcolzMinuteBarWriter,
US_EQUITIES_MINUTES_PER_DAY
)
from pandas.util.testing import assert_series_equal
from six import with_metaclass
from .core import tmp_asset_finder, make_simple_equity_info, gen_calendars
from ..finance.trading import TradingEnvironment
from ..utils import tradingcalendar, factory
from ..utils.classproperty import classproperty
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
from ..utils.metautils import compose_types
from ..pipeline import Pipeline, SimplePipelineEngine
from ..utils.numpy_utils import make_datetime64D
from ..utils.numpy_utils import NaTD
def _take_out_the_trash():
"""Force the gc to clear all innaccessible objects.
This will only kill stranded reference cycles, objects that don't
participate in a cycle are destroyed when there are no more references.
Notes
-----
This function is used to ensure that objects that are holding open file
handles are destroyed, closing the files. On windows files cannot be
deleted if they are opened.
"""
while gc.collect():
pass
class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
@@ -46,6 +82,10 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
@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
@@ -79,7 +119,14 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
@final
@classmethod
def tearDownClass(cls):
_take_out_the_trash()
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
@@ -115,6 +162,7 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
@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
@@ -136,7 +184,10 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
@final
def tearDown(self):
_take_out_the_trash()
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):
@@ -158,6 +209,59 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
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
@@ -166,38 +270,62 @@ class WithLogger(object):
After init_instance_fixtures has been called `self.log_handler` will be a
new ``logbook.NullHandler``.
This behavior may be overridden by defining a ``make_log_handler`` class
method which returns a new logbook.LogHandler instance.
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(object):
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. The default finder is the result of calling
`tmp_asset_finder` with arguments generated as follows::
with an AssetFinder.
equities=cls.make_equities_info(),
futures=cls.make_futures_info(),
exchanges=cls.make_exchanges_info(),
root_symbols=cls.make_root_symbols_info(),
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``.
Each of these methods may be overridden with a function returning a
alternative dataframe of data to write.
The top-level creation behavior can be altered by overriding
`make_asset_finder` as a class method.
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
--------
@@ -207,36 +335,42 @@ class WithAssetFinder(object):
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_equities_info = _make_info
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.equities_info,
futures=cls.futures_info,
exchanges=cls.exchanges_info,
root_symbols=cls.root_symbols_info,
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()
# TODO: Move this to consumers that actually depend on it.
# These are misleading if make_asset_finder is overridden.
cls.equities_info = cls.make_equities_info()
cls.futures_info = cls.make_futures_info()
cls.exchanges_info = cls.make_exchanges_info()
cls.root_symbols_info = cls.make_root_symbols_info()
cls.asset_finder = cls.make_asset_finder()
@@ -248,11 +382,30 @@ class WithTradingEnvironment(WithAssetFinder):
with a trading environment whose `asset_finder` is the result of
`cls.make_asset_finder`.
The ``load`` function may be provided by overriding the
``make_load_function`` class method.
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.
This behavior can be altered by overriding `make_trading_environment` as a
class method.
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
@@ -286,28 +439,52 @@ class WithSimParams(WithTradingEnvironment):
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_START = pd.Timestamp('2006-01-01')
SIM_PARAMS_END = pd.Timestamp('2006-12-31')
SIM_PARAMS_CAPITAL_BASE = float("1.0e5")
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 init_class_fixtures(cls):
super(WithSimParams, cls).init_class_fixtures()
cls.sim_params = factory.create_simulation_parameters(
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):
"""
@@ -318,100 +495,358 @@ class WithNYSETradingDays(object):
(DATA_MAX_DAY - (cls.TRADING_DAY_COUNT) -> DATA_MAX_DAY)
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.
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_MAX_DAY = pd.Timestamp('2016-01-04')
TRADING_DAY_COUNT = 126
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
end_loc = all_days.get_loc(cls.DATA_MAX_DAY)
start_loc = end_loc - cls.TRADING_DAY_COUNT
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 WithTempdir(object):
class WithTmpDir(object):
"""
ZiplineTestCase mixin providing cls.tempdir as a class-level fixture.
ZiplineTestCase mixing providing cls.tmpdir as a class-level fixture.
After init_class_fixtures has been called, `cls.tempdir` is populated
with a TempDirectory object.
After init_class_fixtures has been called, `cls.tmpdir` is populated with
a `testfixtures.TempDirectory` object whose path is `cls.TMP_DIR_PATH`.
The default value of TEMPDIR_PATH is None.
Inheritors can override TEMPDIR_PATH to path argument to `TempDirectory`
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.
"""
TEMPDIR_PATH = None
TMP_DIR_PATH = None
@classmethod
def init_class_fixtures(cls):
super(WithTempdir, cls).init_class_fixtures()
cls.tempdir = TempDirectory(path=cls.TEMPDIR_PATH)
cls.add_class_callback(cls.tempdir.cleanup)
class WithBcolzMinutes(WithTempdir,
WithTradingEnvironment):
"""
ZiplineTestCase mixin providing cls.boclz_minute_bar_reader as a
class-level fixture.
After init_class_fixtures has been called, `cls.bcolz_minute_bar_reader`
is populated with `BcolzMinuteBarReader` with data defined by
`make_bcolz_minute_bar_data`
The default value of BCOLZ_MINUTES_PER_DAY is US_EQUITIES_MINUTES_PER_DAY.
Inheritors can override BCOLZ_MINUTES_PER_DAY to write data for assets
that trade over a different period length.
"""
BCOLZ_MINUTES_PER_DAY = US_EQUITIES_MINUTES_PER_DAY
@classmethod
def init_class_fixtures(cls):
super(WithBcolzMinutes, cls).init_class_fixtures()
writer = BcolzMinuteBarWriter(
cls.env.first_trading_day,
cls.tempdir.path,
cls.env.open_and_closes.market_open,
cls.env.open_and_closes.market_close,
cls.BCOLZ_MINUTES_PER_DAY,
super(WithTmpDir, cls).init_class_fixtures()
cls.tmpdir = cls.enter_class_context(
tmp_dir(path=cls.TMP_DIR_PATH),
)
for sid, data in iteritems(cls.make_bcolz_minute_bar_data()):
writer.write(sid, data)
cls.bcolz_minute_bar_reader = BcolzMinuteBarReader(cls.tempdir.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``.
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')
# allows WithBcolzDailyBarReaderFromCSVs to call the `write_csvs` method
# without needing to reimplement `init_class_fixtures`
_write_method_name = 'write'
@classmethod
def make_bcolz_minute_bar_data(cls):
"""
Returns
-------
A dict of sid -> DataFrame with columns ('open', 'high', 'low',
'close', 'volume') and an index of the minutes on which the prices
occurred.
def make_daily_bar_data(cls):
return create_daily_bar_data(
cls.bcolz_daily_bar_days,
cls.asset_finder.sids,
)
See, zipline.data.minute_bars.BcolzMinuteBarWriter
"""
return {}
@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())
cls.bcolz_daily_bar_reader = BcolzDailyBarReader(t)
class WithPipelineEventDataLoader(WithAssetFinder):
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.
`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.
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):
@@ -421,12 +856,12 @@ class WithPipelineEventDataLoader(WithAssetFinder):
def get_dataset(cls):
return {sid: pd.DataFrame() for sid in cls.get_sids()}
@classmethod
@abstractproperty
def loader_type(self):
return None
raise NotImplementedError('loader_type')
@classmethod
def make_equities_info(cls):
def make_equity_info(cls):
return make_simple_equity_info(
cls.get_sids(),
start_date=pd.Timestamp('2013-01-01', tz='UTC'),
@@ -520,3 +955,74 @@ class WithPipelineEventDataLoader(WithAssetFinder):
assert_series_equal(result[col_name].xs(sid, level=1),
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(),
)
+140
View File
@@ -0,0 +1,140 @@
from six import iteritems, viewkeys, PY2
from zipline.dispatch import dispatch
from zipline.utils.functional import dzip_exact
def _s(word, seq, suffix='s'):
"""Adds a suffix to ``word`` if some sequence has anything other than
exactly one element.
word : str
The string to add the suffix to.
seq : sequence
The sequence to check the length of.
suffix : str, optional.
The suffix to add to ``word``
Returns
-------
maybe_plural : str
``word`` with ``suffix`` added if ``len(seq) != 1``.
"""
return word + (suffix if len(seq) != 1 else '')
def _fmt_path(path):
"""Format the path for final display.
Parameters
----------
path : iterable of str
The path to the values that are not equal.
Returns
-------
fmtd : str
The formatted path to put into the error message.
"""
if not path:
return ''
return 'path: _' + ''.join(path)
@dispatch(object, object)
def assert_equal(result, expected, path=(), **kwargs):
"""Assert that two objects are equal using the ``==`` operator.
Parameters
----------
result : object
The result that came from the function under test.
expected : object
The expected result.
Raises
------
AssertionError
Raised when ``result`` is not equal to ``expected``.
"""
assert result == expected, '%s != %s\n%s' % (
result,
expected,
_fmt_path(path),
)
@assert_equal.register(dict, dict)
def assert_dict_equal(result, expected, path=(), **kwargs):
if path is None:
path = ()
result_keys = viewkeys(result)
expected_keys = viewkeys(expected)
if result_keys != expected_keys:
if result_keys > expected_keys:
diff = result_keys - expected_keys
msg = 'extra %s in result: %r' % (_s('key', diff), diff)
elif result_keys < expected_keys:
diff = expected_keys - result_keys
msg = 'result is missing %s: %r' % (_s('key', diff), diff)
else:
sym = result_keys ^ expected_keys
in_result = sym - expected_keys
in_expected = sym - result_keys
msg = '%s only in result: %s\n%s only in expected: %s' % (
_s('key', in_result),
in_result,
_s('key', in_expected),
in_expected,
)
raise AssertionError(
'dict keys do not match\n%s\n%s' % (
msg,
_fmt_path(path + ('.%s()' % ('viewkeys' if PY2 else 'keys'),)),
),
)
failures = []
for k, (resultv, expectedv) in iteritems(dzip_exact(result, expected)):
try:
assert_equal(
resultv,
expectedv,
path=path + ('[%r]' % k,),
**kwargs
)
except AssertionError as e:
failures.append(str(e))
if failures:
raise AssertionError('\n'.join(failures))
@assert_equal.register(list, list) # noqa
def assert_list_equal(result, expected, path=(), **kwargs):
result_len = len(result)
expected_len = len(expected)
assert result_len == expected_len, (
'list lengths do not match: %d != %d\n%s' %
result_len,
expected_len,
_fmt_path(path),
)
for n, (resultv, expectedv) in enumerate(zip(result, expected)):
assert_equal(
resultv,
expectedv,
path=path + ('[%d]' % n,),
**kwargs
)
try:
# pull the dshape cases in
from datashape.util.testing import assert_dshape_equal
except ImportError:
pass
else:
assert_equal.funcs.update(assert_dshape_equal.funcs)
+1 -1
View File
@@ -1,5 +1,5 @@
"""
Cached object with an expiration date.
Caching utilities for zipline
"""
from collections import namedtuple
+8
View File
@@ -0,0 +1,8 @@
class classproperty(object):
"""Class property
"""
def __init__(self, fget):
self.fget = fget
def __get__(self, instance, owner):
return self.fget(owner)
+32
View File
@@ -18,6 +18,7 @@ import os
import argparse
from copy import copy
import click
from six import print_
from six.moves import configparser
import pandas as pd
@@ -32,6 +33,7 @@ except:
import zipline
from zipline.errors import NoSourceError, PipelineDateError
from .context_tricks import CallbackManager
DEFAULTS = {
'data_frequency': 'daily',
@@ -251,3 +253,33 @@ def run_pipeline(print_algo=True, **kwargs):
perf.to_pickle(output_fname)
return perf
def maybe_show_progress(it, show_progress, **kwargs):
"""Optionally show a progress bar for the given iterator.
Parameters
----------
it : iterable
The underlying iterator.
show_progress : bool
Should progress be shown.
**kwargs
Forwarded to the click progress bar.
Returns
-------
itercontext : context manager
A context manager whose enter is the actual iterator to use.
Examples
--------
with maybe_show_progress([1, 2, 3], True) as ns:
for n in ns:
...
"""
if show_progress:
return click.progressbar(it, **kwargs)
# context manager that just return `it` when we enter it
return CallbackManager(lambda it=it: it)
+19
View File
@@ -0,0 +1,19 @@
from six import PY2
if PY2:
from functools32 import lru_cache
from ctypes import py_object, pythonapi
mappingproxy = pythonapi.PyDictProxy_New
mappingproxy.argtypes = [py_object]
mappingproxy.restype = py_object
else:
from functools import lru_cache
from types import MappingProxyType as mappingproxy
__all__ = [
'lru_cache',
'mappingproxy',
]
+8
View File
@@ -58,6 +58,14 @@ class CallbackManager(object):
def __call__(self, *args, **kwargs):
return _ManagedCallbackContext(self.pre, self.post, args, kwargs)
# special case, if no extra args are passed make this a context manager
# which forwards no args to pre and post
def __enter__(self):
return self.pre()
def __exit__(self, *excinfo):
self.post()
class _ManagedCallbackContext(object):
def __init__(self, pre, post, args, kwargs):
+10
View File
@@ -33,3 +33,13 @@ def invert(d):
except KeyError:
out[v] = {k}
return out
def invert_unique(d, check=True):
"""
Invert a dictionary with unique values into a dictionary with (k, v) pairs
flipped.
"""
if check:
assert len(set(d.values())) == len(d), "Values were not unique!"
return {v: k for k, v in iteritems(d)}
-88
View File
@@ -12,20 +12,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import bisect
import datetime
from collections import MutableMapping
from copy import deepcopy
try:
from six.moves._thread import get_ident
except ImportError:
from six.moves._dummy_thread import get_ident
import numpy as np
import pandas as pd
from toolz import merge
def _ensure_index(x):
@@ -399,82 +390,3 @@ class MutableIndexRollingPanel(object):
self.buffer.loc[non_nan_items, :, non_nan_cols])
self.buffer = new_buffer
class SortedDict(MutableMapping):
"""A mapping of key-value pairs sorted by key according to the sort_key
function provided to the mapping. Ties from the sort_key are broken by
comparing the original keys. `iter` traverses the keys in sort order.
Parameters
----------
key : callable
Called on keys in the mapping to produce the values by which those keys
are sorted.
mapping : mapping, optional
**kwargs
The initial mapping.
>>> d = SortedDict(abs)
>>> d[-1] = 'negative one'
>>> d[0] = 'zero'
>>> d[2] = 'two'
>>> d # doctest: +NORMALIZE_WHITESPACE
SortedDict(<built-in function abs>,
[(0, 'zero'), (-1, 'negative one'), (2, 'two')])
>>> d[1] = 'one' # Mutating the mapping maintains the sort order.
>>> d # doctest: +NORMALIZE_WHITESPACE
SortedDict(<built-in function abs>,
[(0, 'zero'), (-1, 'negative one'), (1, 'one'), (2, 'two')])
>>> del d[0]
>>> d # doctest: +NORMALIZE_WHITESPACE
SortedDict(<built-in function abs>,
[(-1, 'negative one'), (1, 'one'), (2, 'two')])
>>> del d[2]
>>> d
SortedDict(<built-in function abs>, [(-1, 'negative one'), (1, 'one')])
"""
def __init__(self, key, mapping=None, **kwargs):
self._map = {}
self._sorted_key_names = []
self._sort_key = key
self.update(merge(mapping or {}, kwargs))
def __getitem__(self, name):
return self._map[name]
def __setitem__(self, name, value, _bisect_right=bisect.bisect_right):
self._map[name] = value
if len(self._map) > len(self._sorted_key_names):
key = self._sort_key(name)
pair = (key, name)
idx = _bisect_right(self._sorted_key_names, pair)
self._sorted_key_names.insert(idx, pair)
def __delitem__(self, name, _bisect_left=bisect.bisect_left):
del self._map[name]
idx = _bisect_left(self._sorted_key_names,
(self._sort_key(name), name))
del self._sorted_key_names[idx]
def __iter__(self):
for key, name in self._sorted_key_names:
yield name
def __len__(self):
return len(self._map)
def __repr__(self, _repr_running={}):
# Based on OrderedDict/defaultdict
call_key = id(self), get_ident()
if call_key in _repr_running:
return '...'
_repr_running[call_key] = 1
try:
if not self:
return '%s(%r)' % (self.__class__.__name__, self._sort_key)
return '%s(%r, %r)' % (self.__class__.__name__, self._sort_key,
list(self.items()))
finally:
del _repr_running[call_key]
+4 -4
View File
@@ -31,12 +31,12 @@ from zipline.finance.trading import (
SimulationParameters, TradingEnvironment, noop_load
)
from zipline.sources.test_source import create_trade
from zipline.data.loader import ( # For backwards compatibility
load_from_yahoo,
load_bars_from_yahoo,
)
# For backwards compatibility
from zipline.data.loader import (load_from_yahoo,
load_bars_from_yahoo)
__all__ = ['load_from_yahoo', 'load_bars_from_yahoo']
+26 -41
View File
@@ -1,8 +1,6 @@
from abc import ABCMeta, abstractmethod
from weakref import WeakKeyDictionary
from six import with_metaclass, iteritems
from toolz import memoize
# Consistent error to be thrown in various cases regarding overriding
# `final` attributes.
@@ -30,51 +28,38 @@ def is_final(name, mro):
for c in bases_mro(mro))
@memoize(cache=WeakKeyDictionary())
def final_meta_factory(base):
class FinalMeta(type):
"""A metaclass template for classes the want to prevent subclassess from
overriding a some methods or attributes.
"""
Creates a metaclass that inherits from `base` that also checks for `final`
attributes.
This will cause class construction to fail if the class attempts to
override a final method or attribute.
"""
class _FinalMeta(base):
def __new__(mcls, name, bases, dict_):
for k, v in iteritems(dict_):
if is_final(k, bases):
raise _type_error
setattr_ = dict_.get('__setattr__')
if setattr_ is None:
# No `__setattr__` was explicitly defined, look up the super
# class's. `bases[0]` will have a `__setattr__` because
# `object` does so we don't need to worry about the mro.
setattr_ = bases[0].__setattr__
if not is_final('__setattr__', bases) \
and not isinstance(setattr_, final):
# implicitly make the `__setattr__` a `final` object so that
# users cannot just avoid the descriptor protocol.
dict_['__setattr__'] = final(setattr_)
return base.__new__(mcls, name, bases, dict_)
def __setattr__(self, name, value):
"""
This stops the `final` attributes from being reassigned on the
class object.
"""
if is_final(name, self.__mro__):
def __new__(mcls, name, bases, dict_):
for k, v in iteritems(dict_):
if is_final(k, bases):
raise _type_error
base.__setattr__(self, name, value)
setattr_ = dict_.get('__setattr__')
if setattr_ is None:
# No `__setattr__` was explicitly defined, look up the super
# class's. `bases[0]` will have a `__setattr__` because
# `object` does so we don't need to worry about the mro.
setattr_ = bases[0].__setattr__
_FinalMeta.__name__ = '%sFinalMeta' % base.__name__
return _FinalMeta
if not is_final('__setattr__', bases) \
and not isinstance(setattr_, final):
# implicitly make the `__setattr__` a `final` object so that
# users cannot just avoid the descriptor protocol.
dict_['__setattr__'] = final(setattr_)
return super(FinalMeta, mcls).__new__(mcls, name, bases, dict_)
FinalMeta = final_meta_factory(type)
def __setattr__(self, name, value):
"""This stops the `final` attributes from being reassigned on the
class object.
"""
if is_final(name, self.__mro__):
raise _type_error
super(FinalMeta, self).__setattr__(name, value)
class final(with_metaclass(ABCMeta)):
+160 -1
View File
@@ -1,7 +1,60 @@
from pprint import pformat
from six import viewkeys
from six.moves import map
from six.moves import map, zip
from toolz import curry
@curry
def apply(f, *args, **kwargs):
"""Apply a function to arguments.
Parameters
----------
f : callable
The function to call.
*args, **kwargs
**kwargs
Arguments to feed to the callable.
Returns
-------
a : any
The result of ``f(*args, **kwargs)``
Examples
--------
>>> from toolz.curried.operator import add, sub
>>> fs = add(1), sub(1)
>>> tuple(map(apply, fs, (1, 2)))
(2, -1)
Class decorator
>>> instance = apply
>>> @instance
... class obj:
... def f(self):
... return 'f'
...
>>> obj.f()
'f'
>>> issubclass(obj, object)
Traceback (most recent call last):
...
TypeError: issubclass() arg 1 must be a class
>>> isinstance(obj, type)
False
See Also
--------
unpack_apply
mapply
"""
return f(*args, **kwargs)
# Alias for use as a class decorator.
instance = apply
def mapall(funcs, seq):
@@ -83,3 +136,109 @@ def dzip_exact(*dicts):
"dict keys not all equal:\n\n%s" % _format_unequal_keys(dicts)
)
return {k: tuple(d[k] for d in dicts) for k in dicts[0]}
def _gen_unzip(it, elem_len):
"""Helper for unzip which checks the lengths of each element in it.
Parameters
----------
it : iterable[tuple]
An iterable of tuples. ``unzip`` should map ensure that these are
already tuples.
elem_len : int or None
The expected element length. If this is None it is infered from the
length of the first element.
Yields
------
elem : tuple
Each element of ``it``.
Raises
------
ValueError
Raised when the lengths do not match the ``elem_len``.
"""
elem = next(it)
first_elem_len = len(elem)
if elem_len is not None and elem_len != first_elem_len:
raise ValueError(
'element at index 0 was length %d, expected %d' % (
first_elem_len,
elem_len,
)
)
else:
elem_len = first_elem_len
yield elem
for n, elem in enumerate(it, 1):
if len(elem) != elem_len:
raise ValueError(
'element at index %d was length %d, expected %d' % (
n,
len(elem),
elem_len,
),
)
yield elem
def unzip(seq, elem_len=None):
"""Unzip a length n sequence of length m sequences into m seperate length
n sequences.
Parameters
----------
seq : iterable[iterable]
The sequence to unzip.
elem_len : int, optional
The expected length of each element of ``seq``. If not provided this
will be infered from the length of the first element of ``seq``. This
can be used to ensure that code like: ``a, b = unzip(seq)`` does not
fail even when ``seq`` is empty.
Returns
-------
seqs : iterable[iterable]
The new sequences pulled out of the first iterable.
Raises
------
ValueError
Raised when ``seq`` is empty and ``elem_len`` is not provided.
Raised when elements of ``seq`` do not match the given ``elem_len`` or
the length of the first element of ``seq``.
Examples
--------
>>> seq = [('a', 1), ('b', 2), ('c', 3)]
>>> cs, ns = unzip(seq)
>>> cs
('a', 'b', 'c')
>>> ns
(1, 2, 3)
# checks that the elements are the same length
>>> seq = [('a', 1), ('b', 2), ('c', 3, 'extra')]
>>> cs, ns = unzip(seq)
Traceback (most recent call last):
...
ValueError: element at index 2 was length 3, expected 2
# allows an explicit element length instead of infering
>>> seq = [('a', 1, 'extra'), ('b', 2), ('c', 3)]
>>> cs, ns = unzip(seq, 2)
Traceback (most recent call last):
...
ValueError: element at index 0 was length 3, expected 2
# handles empty sequences when a length is given
>>> cs, ns = unzip([], elem_len=2)
>>> cs == ns == ()
True
Notes
-----
This function will force ``seq`` to completion.
"""
ret = tuple(zip(*_gen_unzip(map(tuple, seq), elem_len)))
if ret:
return ret
if elem_len is None:
raise ValueError("cannot unzip empty sequence without 'elem_len'")
return ((),) * elem_len
+4 -1
View File
@@ -71,7 +71,10 @@ def ensure_upper_case(func, argname, arg):
raise TypeError(
"{0}() expected argument '{1}' to"
" be a string, but got {2} instead.".format(
func.__name__, argname, arg,)
func.__name__,
argname,
arg,
),
)
+74
View File
@@ -0,0 +1,74 @@
from operator import attrgetter
def compose_types(a, b, *cs):
"""Compose multiple classes together.
Parameters
----------
*mcls : tuple[type]
The classes that you would like to compose
Returns
-------
cls : type
A type that subclasses all of the types in ``mcls``.
Notes
-----
A common use case for this is to build composed metaclasses, for example,
imagine you have some simple metaclass ``M`` and some instance of ``M``
named ``C`` like so:
.. code-block:: python
class M(type):
def __new__(mcls, name, bases, dict_):
dict_['ayy'] = 'lmao'
return super().__new__(mcls, name, bases, dict_)
class C(metaclass=M):
pass
We now want to create a sublclass of ``C`` that is also an abstract class.
We can use ``compose_types`` to create a new metaclass that is a subclass
of ``M`` and ``ABCMeta``. This is needed because a subclass of a class
with a metaclass must have a metaclass which is a subclass of the metaclass
of the superclass.
.. code-block:: python
class D(C, metaclass=compose_types(M, ABCMeta)):
@abstractmethod
def f(self):
raise NotImplementedError('f')
We can see that this class has both metaclasses applied to it:
.. code-block:: python
>>> D.ayy
lmao
>>> D()
TypeError: Can't instantiate abstract class D with abstract methods f
An important note here is that ``M`` did not use ``type.__new__`` and
instead used ``super()``. This is to support cooperative multiple
inheritence which is needed for ``compose_types`` to work as intended.
After we have composed these types ``M.__new__``\'s super will actually
go to ``ABCMeta.__new__`` and not ``type.__new__``.
Always using ``super()`` to dispatch to your superclass is best practices
anyways so most classes should compose without much special considerations.
"""
mcls = (a, b) + cs
return type(
'compose_types(%s)' % ', '.join(map(attrgetter('__name__'), mcls)),
mcls,
{},
)
+26 -19
View File
@@ -84,26 +84,19 @@ def preprocess(*_unused, **processors):
if defaults is None:
defaults = ()
no_defaults = (NO_DEFAULT,) * (len(args) - len(defaults))
args_defaults = zip(args, no_defaults + defaults)
argset = set(args)
# These assumptions simplify the implementation significantly. If you
# really want to validate a *args/**kwargs function, you'll have to
# implement this here or do it yourself.
args_defaults = list(zip(args, no_defaults + defaults))
if varargs:
raise TypeError(
"Can't validate functions that take *args: %s" % argspec
)
args_defaults.append((varargs, NO_DEFAULT))
if varkw:
raise TypeError(
"Can't validate functions that take **kwargs: %s" % argspec
)
args_defaults.append((varkw, NO_DEFAULT))
argset = set(args) | {varargs, varkw} - {None}
# Arguments can be declared as tuples in Python 2.
if not all(isinstance(arg, str) for arg in args):
raise TypeError(
"Can't validate functions using tuple unpacking: %s" % argspec
"Can't validate functions using tuple unpacking: %s" %
(argspec,)
)
# Ensure that all processors map to valid names.
@@ -113,7 +106,9 @@ def preprocess(*_unused, **processors):
"Got processors for unknown arguments: %s." % bad_names
)
return _build_preprocessed_function(f, processors, args_defaults)
return _build_preprocessed_function(
f, processors, args_defaults, varargs, varkw,
)
return _decorator
@@ -144,7 +139,11 @@ def call(f):
return processor
def _build_preprocessed_function(func, processors, args_defaults):
def _build_preprocessed_function(func,
processors,
args_defaults,
varargs,
varkw):
"""
Build a preprocessed function with the same signature as `func`.
@@ -172,13 +171,21 @@ def _build_preprocessed_function(func, processors, args_defaults):
signature = []
call_args = []
assignments = []
star_map = {
varargs: '*',
varkw: '**',
}
def name_as_arg(arg):
return star_map.get(arg, '') + arg
for arg, default in args_defaults:
if default is NO_DEFAULT:
signature.append(arg)
signature.append(name_as_arg(arg))
else:
default_name = default_name_template % defaults_seen
exec_globals[default_name] = default
signature.append('='.join([arg, default_name]))
signature.append('='.join([name_as_arg(arg), default_name]))
defaults_seen += 1
if arg in processors:
@@ -186,7 +193,7 @@ def _build_preprocessed_function(func, processors, args_defaults):
exec_globals[procname] = processors[arg]
assignments.append(make_processor_assignment(arg, procname))
call_args.append(arg + '=' + arg)
call_args.append(name_as_arg(arg))
exec_str = dedent(
"""\