mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-20 12:20:29 +08:00
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:
@@ -66,3 +66,5 @@ zipline.iml
|
||||
|
||||
# PyCharm custom settings
|
||||
.idea
|
||||
|
||||
TAGS
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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')
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
+61
-133
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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),
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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
@@ -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
@@ -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
File diff suppressed because it is too large
Load Diff
+161
-158
@@ -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
@@ -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():
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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'],
|
||||
)
|
||||
|
||||
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
@@ -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
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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')
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
+623
-117
@@ -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(),
|
||||
)
|
||||
|
||||
@@ -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,5 +1,5 @@
|
||||
"""
|
||||
Cached object with an expiration date.
|
||||
Caching utilities for zipline
|
||||
"""
|
||||
from collections import namedtuple
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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',
|
||||
]
|
||||
@@ -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):
|
||||
|
||||
@@ -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)}
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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(
|
||||
"""\
|
||||
|
||||
Reference in New Issue
Block a user