diff --git a/.gitignore b/.gitignore index 4391b9c9..f47b8f1c 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,5 @@ zipline.iml # PyCharm custom settings .idea + +TAGS diff --git a/docs/source/whatsnew/1.0.0.txt b/docs/source/whatsnew/1.0.0.txt new file mode 100644 index 00000000..eb9565d5 --- /dev/null +++ b/docs/source/whatsnew/1.0.0.txt @@ -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 diff --git a/etc/requirements.txt b/etc/requirements.txt index 6abd8169..a211f690 100644 --- a/etc/requirements.txt +++ b/etc/requirements.txt @@ -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 diff --git a/tests/data/test_us_equity_pricing.py b/tests/data/test_us_equity_pricing.py index c5abfa90..b289c494 100644 --- a/tests/data/test_us_equity_pricing.py +++ b/tests/data/test_us_equity_pricing.py @@ -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 diff --git a/tests/finance/test_slippage.py b/tests/finance/test_slippage.py index b4620003..67f8a37c 100644 --- a/tests/finance/test_slippage.py +++ b/tests/finance/test_slippage.py @@ -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() diff --git a/tests/pipeline/base.py b/tests/pipeline/base.py index 1f89e56f..47f99384 100644 --- a/tests/pipeline/base.py +++ b/tests/pipeline/base.py @@ -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, ) diff --git a/tests/pipeline/test_blaze.py b/tests/pipeline/test_blaze.py index c5e5fa88..fdfe495a 100644 --- a/tests/pipeline/test_blaze.py +++ b/tests/pipeline/test_blaze.py @@ -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') diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index 43f502c2..23e92ae1 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -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 diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 0f3be397..954d4d9a 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -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 diff --git a/tests/pipeline/test_pipeline_algo.py b/tests/pipeline/test_pipeline_algo.py index df0ca7a7..c0b5de80 100644 --- a/tests/pipeline/test_pipeline_algo.py +++ b/tests/pipeline/test_pipeline_algo.py @@ -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 diff --git a/tests/pipeline/test_us_equity_pricing_loader.py b/tests/pipeline/test_us_equity_pricing_loader.py index 19126b35..12a44f95 100644 --- a/tests/pipeline/test_us_equity_pricing_loader.py +++ b/tests/pipeline/test_us_equity_pricing_loader.py @@ -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', ) diff --git a/tests/resources/fetcher_inputs/fetcher_test_data.py b/tests/resources/fetcher_inputs/fetcher_test_data.py index f07f3dfb..fb5e9e1f 100644 --- a/tests/resources/fetcher_inputs/fetcher_test_data.py +++ b/tests/resources/fetcher_inputs/fetcher_test_data.py @@ -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 diff --git a/tests/resources/quandl_samples/AAPL.csv.gz b/tests/resources/quandl_samples/AAPL.csv.gz new file mode 100644 index 00000000..4c2c4e23 Binary files /dev/null and b/tests/resources/quandl_samples/AAPL.csv.gz differ diff --git a/tests/resources/quandl_samples/BRK_A.csv.gz b/tests/resources/quandl_samples/BRK_A.csv.gz new file mode 100644 index 00000000..9a79fe1e Binary files /dev/null and b/tests/resources/quandl_samples/BRK_A.csv.gz differ diff --git a/tests/resources/quandl_samples/MSFT.csv.gz b/tests/resources/quandl_samples/MSFT.csv.gz new file mode 100644 index 00000000..b55a4ec2 Binary files /dev/null and b/tests/resources/quandl_samples/MSFT.csv.gz differ diff --git a/tests/resources/quandl_samples/ZEN.csv.gz b/tests/resources/quandl_samples/ZEN.csv.gz new file mode 100644 index 00000000..1f5f493b Binary files /dev/null and b/tests/resources/quandl_samples/ZEN.csv.gz differ diff --git a/tests/resources/quandl_samples/rebuild_samples.py b/tests/resources/quandl_samples/rebuild_samples.py new file mode 100644 index 00000000..6135ff52 --- /dev/null +++ b/tests/resources/quandl_samples/rebuild_samples.py @@ -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() diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 452d52fb..555e900d 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -15,6 +15,8 @@ from collections import namedtuple import datetime from datetime import timedelta +from textwrap import dedent +from unittest import TestCase, skip import logbook from logbook import TestHandler, WARNING @@ -23,38 +25,29 @@ from nose_parameterized import parameterized from six import iteritems, itervalues from six.moves import range from testfixtures import TempDirectory -from textwrap import dedent -from unittest import TestCase, skip import numpy as np import pandas as pd -from contextlib2 import ExitStack +import pytz +from toolz import merge from zipline import TradingAlgorithm from zipline.api import FixedSlippage +from zipline.assets import Equity, Future +from zipline.assets.synthetic import ( + make_jagged_equity_info, + make_simple_equity_info, +) from zipline.data.data_portal import DataPortal -from zipline.data.minute_bars import BcolzMinuteBarWriter, \ - US_EQUITIES_MINUTES_PER_DAY, BcolzMinuteBarReader +from zipline.data.minute_bars import ( + BcolzMinuteBarReader, + BcolzMinuteBarWriter, + US_EQUITIES_MINUTES_PER_DAY, +) from zipline.data.us_equity_pricing import ( BcolzDailyBarReader, - SQLiteAdjustmentWriter, - SQLiteAdjustmentReader + BcolzDailyBarWriter, ) -from zipline.finance.commission import PerShare -from zipline.finance.execution import LimitOrder -from zipline.finance.order import ORDER_STATUS -from zipline.finance.trading import TradingEnvironment, SimulationParameters -from zipline.sources import DataPanelSource -from zipline.testing.core import ( - FakeDataPortal, - make_trade_data_for_asset_info, - create_data_portal, - create_data_portal_from_trade_history, - DailyBarWriterFromDataFrames, - create_daily_df_for_asset, - write_minute_data_for_asset, - MockDailyBarReader, - make_test_handler) from zipline.errors import ( OrderDuringInitialize, RegisterTradingControlPostInit, @@ -67,7 +60,35 @@ from zipline.errors import ( SetCancelPolicyPostInit, UnsupportedCancelPolicy ) -from zipline.assets import Equity, Future + +from zipline.finance.commission import PerShare +from zipline.finance.execution import LimitOrder +from zipline.finance.order import ORDER_STATUS +from zipline.finance.trading import TradingEnvironment, SimulationParameters +from zipline.sources import DataPanelSource +from zipline.testing import ( + FakeDataPortal, + create_daily_df_for_asset, + create_data_portal, + create_data_portal_from_trade_history, + create_minute_df_for_asset, + empty_trading_env, + make_test_handler, + make_trade_data_for_asset_info, + parameter_space, + str_to_seconds, + tmp_trading_env, + to_utc, + trades_by_sid_to_dfs, +) +from zipline.testing.fixtures import ( + WithDataPortal, + WithLogger, + WithSimParams, + WithTradingEnvironment, + WithTmpDir, + ZiplineTestCase, +) from zipline.test_algorithms import ( access_account_in_init, access_portfolio_in_init, @@ -121,15 +142,7 @@ from zipline.test_algorithms import ( bad_type_history_frequency_kwarg, bad_type_current_assets_kwarg, bad_type_current_fields_kwarg, - no_handle_data) - -from zipline.testing import ( - make_jagged_equity_info, - to_utc, - setup_logger, - teardown_logger, - parameter_space, - str_to_seconds + no_handle_data, ) from zipline.utils.api_support import ZiplineAPI, set_algo_instance from zipline.utils.context_tricks import CallbackManager @@ -145,32 +158,8 @@ from zipline.utils.tradingcalendar import trading_day, trading_days _multiprocess_can_split_ = False -class TestRecordAlgorithm(TestCase): - - @classmethod - def setUpClass(cls): - cls.env = TradingEnvironment() - cls.sids = [133] - cls.env.write_data(equities_identifiers=cls.sids) - - cls.sim_params = factory.create_simulation_parameters( - num_days=4, - env=cls.env - ) - - cls.tempdir = TempDirectory() - - cls.data_portal = create_data_portal( - cls.env, - cls.tempdir, - cls.sim_params, - cls.sids - ) - - @classmethod - def tearDownClass(cls): - del cls.env - cls.tempdir.cleanup() +class TestRecordAlgorithm(WithSimParams, WithDataPortal, ZiplineTestCase): + ASSET_FINDER_EQUITY_SIDS = 133, def test_record_incr(self): algo = RecordAlgorithm(sim_params=self.sim_params, env=self.env) @@ -186,73 +175,59 @@ class TestRecordAlgorithm(TestCase): range(1, len(output) + 1)) -class TestMiscellaneousAPI(TestCase): +class TestMiscellaneousAPI(WithLogger, + WithSimParams, + WithDataPortal, + ZiplineTestCase): + SIM_PARAMS_DATA_FREQUENCY = 'minute' + sids = 1, 2 @classmethod - def setUpClass(cls): - cls.sids = [1, 2] - cls.env = TradingEnvironment() - - metadata = {3: {'symbol': 'PLAY', - 'start_date': '2002-01-01', - 'end_date': '2004-01-01'}, - 4: {'symbol': 'PLAY', - 'start_date': '2005-01-01', - 'end_date': '2006-01-01'}} - - futures_metadata = { - 5: { - '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')}, - 6: { - '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')}, - 7: { - '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')}, - 8: { - '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')} - } - cls.env.write_data(equities_identifiers=cls.sids, - equities_data=metadata, - futures_data=futures_metadata) - - setup_logger(cls) - - cls.sim_params = factory.create_simulation_parameters( - num_days=2, - data_frequency='minute', - emission_rate='daily', - env=cls.env, - ) - - cls.temp_dir = TempDirectory() - - cls.data_portal = create_data_portal( - cls.env, - cls.temp_dir, - cls.sim_params, - cls.sids - ) + def make_equity_info(cls): + return pd.concat(( + make_simple_equity_info(cls.sids, '2002-02-1', '2007-01-01'), + pd.DataFrame.from_dict( + {3: {'symbol': 'PLAY', + 'start_date': '2002-01-01', + 'end_date': '2004-01-01'}, + 4: {'symbol': 'PLAY', + 'start_date': '2005-01-01', + 'end_date': '2006-01-01'}}, + orient='index', + ), + )) @classmethod - def tearDownClass(cls): - del cls.env - teardown_logger(cls) - cls.temp_dir.cleanup() + def make_futures_info(cls): + return pd.DataFrame.from_dict( + { + 5: { + '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')}, + 6: { + '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')}, + 7: { + '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')}, + 8: { + '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')} + }, + orient='index', + ) def test_cancel_policy_outside_init(self): code = """ @@ -318,21 +293,21 @@ def initialize(context): def handle_data(context, data): aapl_dt = data.current(sid(1), "last_traded") - assert aapl_dt == get_datetime() + assert_equal(aapl_dt, get_datetime()) """ algo = TradingAlgorithm(script=algo_text, sim_params=self.sim_params, env=self.env) - + algo.namespace['assert_equal'] = self.assertEqual algo.run(self.data_portal) def test_get_environment(self): expected_env = { 'arena': 'backtest', 'data_frequency': 'minute', - 'start': pd.Timestamp('2006-01-03 14:31:00+0000', tz='UTC'), - 'end': pd.Timestamp('2006-01-04 21:00:00+0000', tz='UTC'), + 'start': pd.Timestamp('2006-01-03 14:31:00+0000', tz='utc'), + 'end': pd.Timestamp('2006-12-29 21:00:00+0000', tz='utc'), 'capital_base': 100000.0, 'platform': 'zipline' } @@ -403,12 +378,19 @@ def handle_data(context, data): def test_schedule_function(self): date_rules = DateRuleFactory time_rules = TimeRuleFactory + us_eastern = pytz.timezone('US/Eastern') def incrementer(algo, data): algo.func_called += 1 + curdt = algo.get_datetime().tz_convert(pytz.utc) self.assertEqual( - algo.get_datetime().time(), - datetime.time(hour=14, minute=31), + curdt, + us_eastern.localize( + datetime.datetime.combine( + curdt.date(), + datetime.time(9, 31) + ), + ), ) def initialize(algo): @@ -476,17 +458,17 @@ def handle_data(context, data): ) algo.run(self.data_portal) - self.assertEqual(len(expected_data), 780) + self.assertEqual(len(expected_data), 97530) self.assertEqual(collected_data_pre, expected_data) self.assertEqual(collected_data_post, expected_data) self.assertEqual( len(function_stack), - 780 * 5, + 97530 * 5, 'Incorrect number of functions called: %s != 780' % len(function_stack), ) - expected_functions = [pre, handle_data, f, g, post] * 780 + expected_functions = [pre, handle_data, f, g, post] * 97530 for n, (f, g) in enumerate(zip(function_stack, expected_functions)): self.assertEqual( f, @@ -707,79 +689,66 @@ def handle_data(context, data): for i, date in enumerate(dates) ] ) - env = TradingEnvironment() - env.write_data(equities_df=metadata) - algo = TradingAlgorithm(env=env) + with tmp_trading_env(equities=metadata) as env: + algo = TradingAlgorithm(env=env) - # Set the period end to a date after the period end - # dates for our assets. - algo.sim_params.period_end = pd.Timestamp('2015-01-01', tz='UTC') + # Set the period end to a date after the period end + # dates for our assets. + algo.sim_params.period_end = pd.Timestamp('2015-01-01', tz='UTC') - # With no symbol lookup date set, we will use the period end date - # for the as_of_date, resulting here in the asset with the earlier - # start date being returned. - result = algo.symbol('DUP') - self.assertEqual(result.symbol, 'DUP') - - # By first calling set_symbol_lookup_date, the relevant asset - # should be returned by lookup_symbol - for i, date in enumerate(dates): - algo.set_symbol_lookup_date(date) + # With no symbol lookup date set, we will use the period end date + # for the as_of_date, resulting here in the asset with the earlier + # start date being returned. result = algo.symbol('DUP') self.assertEqual(result.symbol, 'DUP') - self.assertEqual(result.sid, i + 3) - with self.assertRaises(UnsupportedDatetimeFormat): - algo.set_symbol_lookup_date('foobar') + # By first calling set_symbol_lookup_date, the relevant asset + # should be returned by lookup_symbol + for i, date in enumerate(dates): + algo.set_symbol_lookup_date(date) + result = algo.symbol('DUP') + self.assertEqual(result.symbol, 'DUP') + self.assertEqual(result.sid, i + 3) + + with self.assertRaises(UnsupportedDatetimeFormat): + algo.set_symbol_lookup_date('foobar') -class TestTransformAlgorithm(TestCase): +class TestTransformAlgorithm(WithLogger, + WithDataPortal, + WithSimParams, + ZiplineTestCase): + START_DATE = pd.Timestamp('2006-01-03', tz='utc') + END_DATE = pd.Timestamp('2006-01-06', tz='utc') + + sids = ASSET_FINDER_EQUITY_SIDS = [0, 1, 133] @classmethod - def setUpClass(cls): - setup_logger(cls) - cls.env = TradingEnvironment() - cls.sim_params = factory.create_simulation_parameters(num_days=4, - env=cls.env) - cls.sids = [0, 1, 133] - cls.tempdir = TempDirectory() - - futures_metadata = {3: {'multiplier': 10}} - equities_metadata = {} - - cls.futures_env = TradingEnvironment() - cls.futures_env.write_data(futures_data=futures_metadata) - - for sid in cls.sids: - equities_metadata[sid] = { - 'start_date': cls.sim_params.period_start, - 'end_date': cls.sim_params.period_end - } - - cls.env.write_data(equities_data=equities_metadata, - futures_data=futures_metadata) - - trades_by_sid = {} - for sid in cls.sids: - trades_by_sid[sid] = factory.create_trade_history( - sid, - [10.0, 10.0, 11.0, 11.0], - [100, 100, 100, 300], - timedelta(days=1), - cls.sim_params, - cls.env - ) - - cls.data_portal = create_data_portal_from_trade_history(cls.env, - cls.tempdir, - cls.sim_params, - trades_by_sid) + def make_futures_info(cls): + return pd.DataFrame.from_dict({3: {'multiplier': 10}}, 'index') @classmethod - def tearDownClass(cls): - teardown_logger(cls) - del cls.env - cls.tempdir.cleanup() + def make_daily_bar_data(cls): + return trades_by_sid_to_dfs( + { + sid: factory.create_trade_history( + sid, + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + cls.sim_params, + cls.env + ) for sid in cls.sids + }, + index=cls.sim_params.trading_days, + ) + + @classmethod + def init_class_fixtures(cls): + super(TestTransformAlgorithm, cls).init_class_fixtures() + cls.futures_env = cls.enter_class_context( + tmp_trading_env(futures=cls.make_futures_info()), + ) def test_invalid_order_parameters(self): algo = InvalidOrderAlgorithm( @@ -931,74 +900,40 @@ class TestTransformAlgorithm(TestCase): (TestOrderPercentAlgorithm,) ]) def test_minute_data(self, algo_class): - tempdir = TempDirectory() - - try: - env = TradingEnvironment() - + period_start = pd.Timestamp('2002-1-2', tz='UTC') + period_end = pd.Timestamp('2002-1-4', tz='UTC') + equities = pd.DataFrame([{ + 'start_date': period_start, + 'end_date': period_end + timedelta(days=1) + }] * 2) + with TempDirectory() as tempdir, \ + tmp_trading_env(equities=equities) as env: sim_params = SimulationParameters( - period_start=pd.Timestamp('2002-1-2', tz='UTC'), - period_end=pd.Timestamp('2002-1-4', tz='UTC'), + period_start=period_start, + period_end=period_end, capital_base=float("1.0e5"), data_frequency='minute', env=env ) - equities_metadata = {} - - for sid in [0, 1]: - equities_metadata[sid] = { - 'start_date': sim_params.period_start, - 'end_date': sim_params.period_end + timedelta(days=1) - } - - env.write_data(equities_data=equities_metadata) - data_portal = create_data_portal( env, tempdir, sim_params, - [0, 1] + equities.index, ) - algo = algo_class(sim_params=sim_params, env=env) algo.run(data_portal) - finally: - tempdir.cleanup() -class TestPositions(TestCase): - @classmethod - def setUpClass(cls): - setup_logger(cls) - cls.env = TradingEnvironment() - cls.sim_params = factory.create_simulation_parameters(num_days=4, - env=cls.env) +class TestPositions(WithLogger, + WithDataPortal, + WithSimParams, + ZiplineTestCase): + START_DATE = pd.Timestamp('2006-01-03', tz='utc') + END_DATE = pd.Timestamp('2006-01-06', tz='utc') - cls.sids = [1, 133] - cls.tempdir = TempDirectory() - - equities_metadata = {} - - for sid in cls.sids: - equities_metadata[sid] = { - 'start_date': cls.sim_params.period_start, - 'end_date': cls.sim_params.period_end - } - - cls.env.write_data(equities_data=equities_metadata) - - cls.data_portal = create_data_portal( - cls.env, - cls.tempdir, - cls.sim_params, - cls.sids - ) - - @classmethod - def tearDownClass(cls): - teardown_logger(cls) - cls.tempdir.cleanup() + sids = ASSET_FINDER_EQUITY_SIDS = [1, 133] def test_empty_portfolio(self): algo = EmptyPositionsAlgorithm(self.sids, @@ -1028,143 +963,80 @@ class TestPositions(TestCase): self.assertTrue(empty_positions.all()) -class TestBeforeTradingStart(TestCase): +class TestBeforeTradingStart(WithDataPortal, + WithSimParams, + ZiplineTestCase): + START_DATE = pd.Timestamp('2016-01-06', tz='utc') + END_DATE = pd.Timestamp('2016-01-07', tz='utc') + SIM_PARAMS_CAPITAL_BASE = 10000 + SIM_PARAMS_DATA_FREQUENCY = 'minute' + BCOLZ_DAILY_BAR_LOOKBACK_DAYS = BCOLZ_MINUTE_BAR_LOOKBACK_DAYS = 1 + + data_start = ASSET_FINDER_EQUITY_START_DATE = pd.Timestamp( + '2016-01-05', + tz='utc', + ) + + SPLIT_ASSET_SID = 3 + ASSET_FINDER_EQUITY_SIDS = 1, 2, SPLIT_ASSET_SID + @classmethod - def setUpClass(cls): - cls.env = TradingEnvironment() - cls.tempdir = TempDirectory() - - cls.trading_days = cls.env.days_in_range( - start=pd.Timestamp("2016-01-05", tz='UTC'), - end=pd.Timestamp("2016-01-07", tz='UTC') - ) - - equities_data = {} - for sid in [1, 2, 3]: - equities_data[sid] = { - "start_date": cls.trading_days[0], - "end_date": cls.trading_days[-1], - "symbol": "ASSET{0}".format(sid), - } - - cls.env.write_data(equities_data=equities_data) - - 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) - - 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, 8554]: - write_minute_data_for_asset( - cls.env, minute_writer, cls.trading_days[0], - cls.trading_days[-1], sid - ) - - # Write data with split asset + def make_minute_bar_data(cls): asset_minutes = cls.env.minutes_for_days_in_range( - cls.trading_days[0], cls.trading_days[-1]) + cls.data_start, + cls.END_DATE, + ) minutes_count = len(asset_minutes) - minutes_arr = np.array(range(1, 1 + minutes_count)) - - df = pd.DataFrame({ - "open": minutes_arr + 1, - "high": minutes_arr + 2, - "low": minutes_arr - 1, - "close": minutes_arr, - "volume": 100 * minutes_arr, - "dt": asset_minutes - }).set_index("dt") - df.iloc[780:] = df.iloc[780:] / 2.0 - - minute_writer.write(3, df) - - # asset2 only trades every 50 minutes - write_minute_data_for_asset( - cls.env, minute_writer, cls.trading_days[0], - cls.trading_days[-1], 2, 50 + minutes_arr = np.arange(minutes_count) + 1 + split_data = pd.DataFrame( + { + 'open': minutes_arr + 1, + 'high': minutes_arr + 2, + 'low': minutes_arr - 1, + 'close': minutes_arr, + 'volume': 100 * minutes_arr, + }, + index=asset_minutes, ) - - minute_reader = BcolzMinuteBarReader(cls.tempdir.path) - adj_reader = cls.create_adjustments_reader() - - daily_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(daily_path, cls.trading_days, dfs) - - cls.sim_params = SimulationParameters( - period_start=cls.trading_days[1], - period_end=cls.trading_days[-1], - data_frequency="minute", - env=cls.env - ) - - cls.data_portal = DataPortal( - env=cls.env, - equity_daily_reader=BcolzDailyBarReader(daily_path), - equity_minute_reader=minute_reader, - adjustment_reader=adj_reader + split_data.iloc[780:] = split_data.iloc[780:] / 2.0 + return merge( + { + sid: create_minute_df_for_asset( + cls.env, + cls.data_start, + cls.sim_params.period_end, + ) + for sid in (1, 8554) + }, + { + 2: create_minute_df_for_asset( + cls.env, + cls.data_start, + cls.sim_params.period_end, + 50, + ), + cls.SPLIT_ASSET_SID: split_data, + }, ) @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-07"), + 'effective_date': str_to_seconds('2016-01-07'), 'ratio': 0.5, - 'sid': cls.SPLIT_ASSET.sid + 'sid': cls.SPLIT_ASSET_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) - @classmethod - def tearDownClass(cls): - del cls.data_portal - del cls.env - cls.tempdir.cleanup() + def make_daily_bar_data(cls): + for sid in cls.ASSET_FINDER_EQUITY_SIDS: + yield sid, create_daily_df_for_asset( + cls.env, + cls.data_start, + cls.sim_params.period_end, + ) def test_data_in_bts_minute(self): algo_code = dedent(""" @@ -1191,7 +1063,6 @@ class TestBeforeTradingStart(TestCase): algo = TradingAlgorithm( script=algo_code, - data_frequency="minute", sim_params=self.sim_params, env=self.env ) @@ -1239,7 +1110,7 @@ class TestBeforeTradingStart(TestCase): np.full(40, np.nan), algo.history_values[0]["high"][2][20:] ) - def data_in_bts_daily(self): + def test_data_in_bts_daily(self): algo_code = dedent(""" from zipline.api import record, sid def initialize(context): @@ -1264,7 +1135,6 @@ class TestBeforeTradingStart(TestCase): algo = TradingAlgorithm( script=algo_code, - data_frequency="minute", sim_params=self.sim_params, env=self.env ) @@ -1444,9 +1314,15 @@ class TestBeforeTradingStart(TestCase): 10000 + 780 - 392 - 1) -class TestAlgoScript(TestCase): +class TestAlgoScript(WithLogger, + WithDataPortal, + WithSimParams, + ZiplineTestCase): + START_DATE = pd.Timestamp('2006-01-03', tz='utc') + END_DATE = pd.Timestamp('2006-12-31', tz='utc') + BCOLZ_DAILY_BAR_LOOKBACK_DAYS = 5 # max history window length - ARG_TYPE_TEST_CASES = [ + ARG_TYPE_TEST_CASES = ( ('history__assets', (bad_type_history_assets, 'Asset, str', True)), ('history__fields', (bad_type_history_fields, 'str', True)), ('history__bar_count', (bad_type_history_bar_count, 'int', False)), @@ -1467,61 +1343,42 @@ class TestAlgoScript(TestCase): (bad_type_current_assets_kwarg, 'Asset, str', True)), ('current_kwarg__fields', (bad_type_current_fields_kwarg, 'str', True)), - ] + ) + + sids = 0, 1, 3, 133 @classmethod - def setUpClass(cls): - setup_logger(cls) - cls.env = TradingEnvironment() - cls.sim_params = factory.create_simulation_parameters(num_days=251, - env=cls.env) - - cls.sids = [0, 1, 3, 133] - cls.tempdir = TempDirectory() - - equities_metadata = {} - - for sid in cls.sids: - equities_metadata[sid] = { - 'start_date': cls.sim_params.period_start, - 'end_date': cls.env.next_trading_day(cls.sim_params.period_end) - } - - if sid == 3: - equities_metadata[sid]["symbol"] = "TEST" - equities_metadata[sid]["asset_type"] = "equity" - - cls.env.write_data(equities_data=equities_metadata) - - days = 251 - - cls.trades_by_sid = { - 0: factory.create_trade_history( - 0, - [10.0] * days, - [100] * days, - timedelta(days=1), - cls.sim_params, - cls.env), - 3: factory.create_trade_history( - 3, - [10.0] * days, - [100] * days, - timedelta(days=1), - cls.sim_params, - cls.env) - } - - cls.data_portal = create_data_portal_from_trade_history( - cls.env, cls.tempdir, cls.sim_params, cls.trades_by_sid + def make_equity_info(cls): + data = make_simple_equity_info( + cls.sids, + cls.START_DATE, + cls.END_DATE, ) + data.loc[3, 'symbol'] = 'TEST' + return data @classmethod - def tearDownClass(cls): - del cls.data_portal - del cls.env - cls.tempdir.cleanup() - teardown_logger(cls) + def make_daily_bar_data(cls): + days = len(cls.env.days_in_range(cls.START_DATE, cls.END_DATE)) + return trades_by_sid_to_dfs( + { + 0: factory.create_trade_history( + 0, + [10.0] * days, + [100] * days, + timedelta(days=1), + cls.sim_params, + cls.env), + 3: factory.create_trade_history( + 3, + [10.0] * days, + [100] * days, + timedelta(days=1), + cls.sim_params, + cls.env) + }, + index=cls.sim_params.trading_days, + ) def test_noop(self): algo = TradingAlgorithm(initialize=initialize_noop, @@ -1807,56 +1664,24 @@ def handle_data(context, data): Test that api methods on the data object can be called with positional arguments. """ - data_portal = create_data_portal_from_trade_history( - self.env, - self.tempdir, - self.sim_params, - self.trades_by_sid - ) - - sim_params = SimulationParameters( - period_start=pd.Timestamp('2006-02-01', tz='UTC'), - period_end=pd.Timestamp('2006-03-01', tz='UTC'), - capital_base=self.sim_params.capital_base, - data_frequency=self.sim_params.data_frequency, - emission_rate=self.sim_params.emission_rate, - env=self.env, - ) - test_algo = TradingAlgorithm( script=call_without_kwargs, - sim_params=sim_params, + sim_params=self.sim_params, env=self.env, ) - test_algo.run(data_portal) + test_algo.run(self.data_portal) def test_good_kwargs(self): """ Test that api methods on the data object can be called with keyword arguments. """ - data_portal = create_data_portal_from_trade_history( - self.env, - self.tempdir, - self.sim_params, - self.trades_by_sid - ) - - sim_params = SimulationParameters( - period_start=pd.Timestamp('2006-02-01', tz='UTC'), - period_end=pd.Timestamp('2006-03-01', tz='UTC'), - capital_base=self.sim_params.capital_base, - data_frequency=self.sim_params.data_frequency, - emission_rate=self.sim_params.emission_rate, - env=self.env, - ) - test_algo = TradingAlgorithm( script=call_with_kwargs, - sim_params=sim_params, + sim_params=self.sim_params, env=self.env, ) - test_algo.run(data_portal) + test_algo.run(self.data_portal) @parameterized.expand([('history', call_with_bad_kwargs_history), ('current', call_with_bad_kwargs_current)]) @@ -1866,20 +1691,13 @@ def handle_data(context, data): a meaningful TypeError that we create, rather than an unhelpful cython error """ - data_portal = create_data_portal_from_trade_history( - self.env, - self.tempdir, - self.sim_params, - self.trades_by_sid - ) - with self.assertRaises(TypeError) as cm: test_algo = TradingAlgorithm( script=algo_text, sim_params=self.sim_params, env=self.env, ) - test_algo.run(data_portal) + test_algo.run(self.data_portal) self.assertEqual("%s() got an unexpected keyword argument 'blahblah'" % name, cm.exception.args[0]) @@ -1921,37 +1739,15 @@ def handle_data(context, data): algo.run(self.data_portal) -class TestGetDatetime(TestCase): +class TestGetDatetime(WithLogger, + WithSimParams, + WithDataPortal, + ZiplineTestCase): + SIM_PARAMS_DATA_FREQUENCY = 'minute' + START_DATE = to_utc('2014-01-02 9:31') + END_DATE = to_utc('2014-01-03 9:31') - @classmethod - def setUpClass(cls): - cls.env = TradingEnvironment() - cls.env.write_data(equities_identifiers=[0, 1]) - - setup_logger(cls) - - cls.sim_params = factory.create_simulation_parameters( - data_frequency='minute', - env=cls.env, - start=to_utc('2014-01-02 9:31'), - end=to_utc('2014-01-03 9:31') - ) - - cls.tempdir = TempDirectory() - - cls.data_portal = create_data_portal( - cls.env, - cls.tempdir, - cls.sim_params, - [1] - ) - - @classmethod - def tearDownClass(cls): - del cls.data_portal - del cls.env - teardown_logger(cls) - cls.tempdir.cleanup() + ASSET_FINDER_EQUITY_SIDS = 0, 1 @parameterized.expand( [ @@ -1994,43 +1790,18 @@ class TestGetDatetime(TestCase): self.assertFalse(algo.first_bar) -class TestTradingControls(TestCase): +class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): + START_DATE = pd.Timestamp('2006-01-03', tz='utc') + END_DATE = pd.Timestamp('2006-01-06', tz='utc') + + sid = 133 + sids = ASSET_FINDER_EQUITY_SIDS = 133, 134 @classmethod - def setUpClass(cls): - cls.sid = 133 - cls.env = TradingEnvironment() - cls.sim_params = factory.create_simulation_parameters(num_days=4, - env=cls.env) - - cls.env.write_data(equities_data={ - 133: { - 'start_date': cls.sim_params.period_start, - 'end_date': cls.env.next_trading_day(cls.sim_params.period_end) - }, - 134: { - 'start_date': cls.sim_params.period_start, - 'end_date': cls.env.next_trading_day(cls.sim_params.period_end) - } - }) - - cls.asset = cls.env.asset_finder.retrieve_asset(cls.sid) - cls.another_asset = cls.env.asset_finder.retrieve_asset(134) - - cls.tempdir = TempDirectory() - - cls.data_portal = create_data_portal( - cls.env, - cls.tempdir, - cls.sim_params, - [cls.sid] - ) - - @classmethod - def tearDownClass(cls): - del cls.data_portal - del cls.env - cls.tempdir.cleanup() + def init_class_fixtures(cls): + super(TestTradingControls, cls).init_class_fixtures() + cls.asset = cls.asset_finder.retrieve_asset(cls.sid) + cls.another_asset = cls.asset_finder.retrieve_asset(134) def _check_algo(self, algo, @@ -2149,6 +1920,7 @@ class TestTradingControls(TestCase): def handle_data(algo, data): algo.order(algo.sid(self.sid), 1) algo.order_count += 1 + algo = SetMaxOrderSizeAlgorithm(asset=self.asset, max_shares=10, max_notional=500.0, @@ -2207,18 +1979,24 @@ class TestTradingControls(TestCase): self.check_algo_fails(algo, handle_data, 0) def test_set_max_order_count(self): - tempdir = TempDirectory() - try: - env = TradingEnvironment() - sim_params = factory.create_simulation_parameters( - num_days=4, env=env, data_frequency="minute") - - env.write_data(equities_data={ + start = pd.Timestamp('2006-01-05', tz='utc') + metadata = pd.DataFrame.from_dict( + { 1: { - 'start_date': sim_params.period_start, - 'end_date': sim_params.period_end + timedelta(days=1) - } - }) + 'start_date': start, + 'end_date': start + timedelta(days=6) + }, + }, + orient='index', + ) + with TempDirectory() as tempdir, \ + tmp_trading_env(equities=metadata) as env: + sim_params = factory.create_simulation_parameters( + start=start, + num_days=4, + env=env, + data_frequency='minute', + ) data_portal = create_data_portal( env, @@ -2272,8 +2050,6 @@ class TestTradingControls(TestCase): env=env) algo._handle_data = handle_data3 algo.run(data_portal) - finally: - tempdir.cleanup() def test_long_only(self): # Sell immediately -> fail immediately. @@ -2333,122 +2109,85 @@ class TestTradingControls(TestCase): algo.run(self.data_portal) def test_asset_date_bounds(self): - tempdir = TempDirectory() - try: - # Run the algorithm with a sid that ends far in the future - temp_env = TradingEnvironment() - + metadata = pd.DataFrame([{ + 'start_date': self.sim_params.period_start, + 'end_date': '2020-01-01', + }]) + with TempDirectory() as tempdir, \ + tmp_trading_env(equities=metadata) as env: + algo = SetAssetDateBoundsAlgorithm( + sim_params=self.sim_params, + env=env, + ) data_portal = create_data_portal( - temp_env, + env, tempdir, self.sim_params, [0] ) - - metadata = {0: {'start_date': self.sim_params.period_start, - 'end_date': '2020-01-01'}} - - algo = SetAssetDateBoundsAlgorithm( - equities_metadata=metadata, - sim_params=self.sim_params, - env=temp_env, - ) algo.run(data_portal) - finally: - tempdir.cleanup() - - # Run the algorithm with a sid that has already ended - tempdir = TempDirectory() - try: - temp_env = TradingEnvironment() + metadata = pd.DataFrame([{ + 'start_date': '1989-01-01', + 'end_date': '1990-01-01', + }]) + with TempDirectory() as tempdir, \ + tmp_trading_env(equities=metadata) as env: data_portal = create_data_portal( - temp_env, + env, tempdir, self.sim_params, [0] ) - metadata = {0: {'start_date': '1989-01-01', - 'end_date': '1990-01-01'}} - algo = SetAssetDateBoundsAlgorithm( - equities_metadata=metadata, sim_params=self.sim_params, - env=temp_env, + env=env, ) with self.assertRaises(TradingControlViolation): algo.run(data_portal) - finally: - tempdir.cleanup() - # Run the algorithm with a sid that has not started - tempdir = TempDirectory() - try: - temp_env = TradingEnvironment() + metadata = pd.DataFrame([{ + 'start_date': '2020-01-01', + 'end_date': '2021-01-01', + }]) + with TempDirectory() as tempdir, \ + tmp_trading_env(equities=metadata) as env: data_portal = create_data_portal( - temp_env, + env, tempdir, self.sim_params, [0] ) - - metadata = {0: {'start_date': '2020-01-01', - 'end_date': '2021-01-01'}} - algo = SetAssetDateBoundsAlgorithm( - equities_metadata=metadata, sim_params=self.sim_params, - env=temp_env, + env=env, ) - with self.assertRaises(TradingControlViolation): algo.run(data_portal) - finally: - tempdir.cleanup() +class TestAccountControls(WithDataPortal, WithSimParams, ZiplineTestCase): + START_DATE = pd.Timestamp('2006-01-03', tz='utc') + END_DATE = pd.Timestamp('2006-01-06', tz='utc') -class TestAccountControls(TestCase): + sidint, = ASSET_FINDER_EQUITY_SIDS = (133,) @classmethod - def setUpClass(cls): - cls.sidint = 133 - cls.env = TradingEnvironment() - cls.sim_params = factory.create_simulation_parameters( - num_days=4, env=cls.env + def make_daily_bar_data(cls): + return trades_by_sid_to_dfs( + { + cls.sidint: factory.create_trade_history( + cls.sidint, + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + cls.sim_params, + cls.env, + ), + }, + index=cls.sim_params.trading_days, ) - cls.env.write_data(equities_data={ - 133: { - 'start_date': cls.sim_params.period_start, - 'end_date': cls.sim_params.period_end + timedelta(days=1) - } - }) - - cls.tempdir = TempDirectory() - - trades_by_sid = { - cls.sidint: factory.create_trade_history( - cls.sidint, - [10.0, 10.0, 11.0, 11.0], - [100, 100, 100, 300], - timedelta(days=1), - cls.sim_params, - cls.env, - ) - } - - cls.data_portal = create_data_portal_from_trade_history(cls.env, - cls.tempdir, - cls.sim_params, - trades_by_sid) - - @classmethod - def tearDownClass(cls): - del cls.data_portal - del cls.env - cls.tempdir.cleanup() - def _check_algo(self, algo, handle_data, @@ -2575,46 +2314,26 @@ class TestAccountControls(TestCase): # format(i, actual_position, expected_positions[i])) -class TestFutureFlip(TestCase): - @classmethod - def setUpClass(cls): - cls.tempdir = TempDirectory() - - cls.env = TradingEnvironment() - cls.days = pd.date_range(start=pd.Timestamp("2006-01-09", tz='UTC'), - end=pd.Timestamp("2006-01-12", tz='UTC')) - - cls.sid = 1 - - cls.sim_params = factory.create_simulation_parameters( - start=cls.days[0], - end=cls.days[-2] - ) - - trades = factory.create_trade_history( - cls.sid, - [1, 2, 4], - [1e9, 1e9, 1e9], - timedelta(days=1), - cls.sim_params, - cls.env - ) - - trades_by_sid = { - cls.sid: trades - } - - cls.data_portal = create_data_portal_from_trade_history( - cls.env, - cls.tempdir, - cls.sim_params, - trades_by_sid - ) +class TestFutureFlip(WithSimParams, WithDataPortal, ZiplineTestCase): + START_DATE = pd.Timestamp('2006-01-09', tz='utc') + END_DATE = pd.Timestamp('2006-01-10', tz='utc') + sid, = ASSET_FINDER_EQUITY_SIDS = (1,) @classmethod - def tearDownClass(cls): - del cls.data_portal - cls.tempdir.cleanup() + def make_daily_bar_data(cls): + return trades_by_sid_to_dfs( + { + cls.sid: factory.create_trade_history( + cls.sid, + [1, 2, 4], + [1e9, 1e9, 1e9], + timedelta(days=1), + cls.sim_params, + cls.env + ), + }, + index=cls.sim_params.trading_days, + ) @skip def test_flip_algo(self): @@ -2655,11 +2374,7 @@ class TestFutureFlip(TestCase): format(i, actual_position, expected_positions[i])) -class TestTradingAlgorithm(TestCase): - def setUp(self): - self.env = TradingEnvironment() - self.days = self.env.trading_days[:4] - +class TestTradingAlgorithm(ZiplineTestCase): def test_analyze_called(self): self.perf_ref = None @@ -2672,123 +2387,92 @@ class TestTradingAlgorithm(TestCase): def analyze(context, perf): self.perf_ref = perf - algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data, - analyze=analyze) + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + analyze=analyze, + ) - data_portal = FakeDataPortal(self.env) + with empty_trading_env() as env: + data_portal = FakeDataPortal(env) + results = algo.run(data_portal) - results = algo.run(data_portal) self.assertIs(results, self.perf_ref) -class TestOrderCancelation(TestCase): - @classmethod - def setUpClass(cls): - cls.env = TradingEnvironment() - cls.tempdir = TempDirectory() +class TestOrderCancelation(WithDataPortal, + WithSimParams, + ZiplineTestCase): - cls.days = cls.env.days_in_range( - start=pd.Timestamp("2016-01-05", tz='UTC'), - end=pd.Timestamp("2016-01-07", tz='UTC') + START_DATE = pd.Timestamp('2016-01-05', tz='utc') + END_DATE = pd.Timestamp('2016-01-07', tz='utc') + + ASSET_FINDER_EQUITY_SIDS = (1,) + ASSET_FINDER_EQUITY_SYMBOLS = ('ASSET1',) + + code = dedent( + """ + from zipline.api import ( + sid, order, set_slippage, slippage, VolumeShareSlippage, + set_cancel_policy, cancel_policy, EODCancel ) - cls.env.write_data(equities_data={ - 1: { - 'start_date': cls.days[0], - 'end_date': cls.days[-1], - 'symbol': "ASSET1" - } - }) - cls.data_portal = DataPortal( - cls.env, - equity_minute_reader=cls.build_minute_data(), - equity_daily_reader=cls.build_daily_data() - ) - - cls.code = dedent( - """ - from zipline.api import ( - sid, order, set_slippage, slippage, VolumeShareSlippage, - set_cancel_policy, cancel_policy, EODCancel + def initialize(context): + set_slippage( + slippage.VolumeShareSlippage( + volume_limit=1, + price_impact=0 + ) ) - def initialize(context): - set_slippage( - slippage.VolumeShareSlippage( - volume_limit=1, - price_impact=0 - ) - ) + {0} + context.ordered = False - {0} - context.ordered = False - def handle_data(context, data): - if not context.ordered: - order(sid(1), {1}) - context.ordered = True - """ - ) + def handle_data(context, data): + if not context.ordered: + order(sid(1), {1}) + context.ordered = True + """, + ) @classmethod - def tearDownClass(cls): - del cls.data_portal - cls.tempdir.cleanup() - - @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 - ) - + def make_minute_bar_data(cls): asset_minutes = cls.env.minutes_for_days_in_range( - cls.days[0], cls.days[-1] + cls.sim_params.period_start, + cls.sim_params.period_end, ) minutes_count = len(asset_minutes) - minutes_arr = np.array(range(1, 1 + minutes_count)) + minutes_arr = np.arange(1, 1 + minutes_count) # normal test data, but volume is pinned at 1 share per minute - df = pd.DataFrame({ - "open": minutes_arr + 1, - "high": minutes_arr + 2, - "low": minutes_arr - 1, - "close": minutes_arr, - "volume": np.full(minutes_count, 1), - "dt": asset_minutes - }).set_index("dt") - - writer.write(1, df) - - return BcolzMinuteBarReader(cls.tempdir.path) - - @classmethod - def build_daily_data(cls): - path = cls.tempdir.getpath("testdaily.bcolz") - - dfs = { - 1: pd.DataFrame({ - "open": np.full(3, 1), - "high": np.full(3, 1), - "low": np.full(3, 1), - "close": np.full(3, 1), - "volume": np.full(3, 1), - "day": [day.value for day in cls.days] - }) + return { + 1: pd.DataFrame( + { + 'open': minutes_arr + 1, + 'high': minutes_arr + 2, + 'low': minutes_arr - 1, + 'close': minutes_arr, + 'volume': np.full(minutes_count, 1), + }, + index=asset_minutes, + ), } - daily_writer = DailyBarWriterFromDataFrames(dfs) - daily_writer.write(path, cls.days, dfs) - - return BcolzDailyBarReader(path) + @classmethod + def make_daily_bar_data(cls): + yield 1, pd.DataFrame( + { + 'open': np.full(3, 1), + 'high': np.full(3, 1), + 'low': np.full(3, 1), + 'close': np.full(3, 1), + 'volume': np.full(3, 1), + }, + index=cls.sim_params.trading_days, + ) def prep_algo(self, cancelation_string, data_frequency="minute", amount=1000): @@ -2797,8 +2481,8 @@ class TestOrderCancelation(TestCase): script=code, env=self.env, sim_params=SimulationParameters( - period_start=self.days[0], - period_end=self.days[-1], + period_start=self.sim_params.period_start, + period_end=self.sim_params.period_end, env=self.env, data_frequency=data_frequency ) @@ -2814,7 +2498,7 @@ class TestOrderCancelation(TestCase): # so the order should be cancelled at the end of the day. algo = self.prep_algo( "set_cancel_policy(cancel_policy.EODCancel())", - amount=(direction * 1000) + amount=np.copysign(1000, direction), ) log_catcher = TestHandler() @@ -2823,7 +2507,10 @@ class TestOrderCancelation(TestCase): for daily_positions in results.positions: self.assertEqual(1, len(daily_positions)) - self.assertEqual(direction * 389, daily_positions[0]["amount"]) + self.assertEqual( + np.copysign(389, direction), + daily_positions[0]["amount"], + ) self.assertEqual(1, results.positions[0][0]["sid"].sid) # should be an order on day1, but no more orders afterwards @@ -2837,7 +2524,7 @@ class TestOrderCancelation(TestCase): the_order = results.orders[0][0] self.assertEqual(ORDER_STATUS.CANCELLED, the_order["status"]) - self.assertEqual(direction * 389, the_order["filled"]) + self.assertEqual(np.copysign(389, direction), the_order["filled"]) warnings = [record for record in log_catcher.records if record.level == WARNING] @@ -2971,13 +2658,14 @@ class TestRemoveData(TestCase): self.assertEqual(algo.data_lengths, self.live_asset_counts) -class TestEquityAutoClose(TestCase): +class TestEquityAutoClose(WithTmpDir, ZiplineTestCase): """ Tests if delisted equities are properly removed from a portfolio holding positions in said equities. """ @classmethod - def setUpClass(cls): + def init_class_fixtures(cls): + super(TestEquityAutoClose, cls).init_class_fixtures() start_date = pd.Timestamp('2015-01-05', tz='UTC') start_date_loc = trading_days.get_loc(start_date) test_duration = 7 @@ -2985,20 +2673,6 @@ class TestEquityAutoClose(TestCase): start_date_loc:start_date_loc + test_duration ] cls.first_asset_expiration = cls.test_days[2] - cls.tempdir = TempDirectory() - - def setUp(self): - self._teardown_stack = ExitStack() - - def tearDown(self): - self._teardown_stack.close() - - @classmethod - def tearDownClass(cls): - cls.tempdir.cleanup() - - def make_temp_resource(self, resource_context): - return self._teardown_stack.enter_context(resource_context) def make_data(self, auto_close_delta, frequency, capital_base=float("1.0e5")): @@ -3012,10 +2686,9 @@ class TestEquityAutoClose(TestCase): auto_close_delta=auto_close_delta, ) - sids = asset_info.keys() + sids = asset_info.index - env = TradingEnvironment() - env.write_data(equities_data=asset_info) + env = self.enter_instance_context(tmp_trading_env(equities=asset_info)) market_opens = env.open_and_closes.market_open.loc[self.test_days] market_closes = env.open_and_closes.market_close.loc[self.test_days] @@ -3032,9 +2705,10 @@ class TestEquityAutoClose(TestCase): volume_step_by_date=10, frequency=frequency ) - path = self.tempdir.getpath("testdaily.bcolz") - writer = DailyBarWriterFromDataFrames(trade_data_by_sid) - writer.write(path, dates, trade_data_by_sid) + path = self.tmpdir.getpath("testdaily.bcolz") + BcolzDailyBarWriter(path, dates).write( + iteritems(trade_data_by_sid), + ) data_portal = DataPortal( env, equity_daily_reader=BcolzDailyBarReader(path) @@ -3046,7 +2720,7 @@ class TestEquityAutoClose(TestCase): ) writer = BcolzMinuteBarWriter( self.test_days[0], - self.tempdir.path, + self.tmpdir.path, market_opens, market_closes, US_EQUITIES_MINUTES_PER_DAY @@ -3065,7 +2739,7 @@ class TestEquityAutoClose(TestCase): ) data_portal = DataPortal( env, - equity_minute_reader=BcolzMinuteBarReader(self.tempdir.path) + equity_minute_reader=BcolzMinuteBarReader(self.tmpdir.path) ) else: self.fail("Unknown frequency in make_data: %r" % frequency) @@ -3082,13 +2756,9 @@ class TestEquityAutoClose(TestCase): ) if frequency == 'daily': - trade_data_by_sid = { - sid: df.set_index('day') for sid, df in - iteritems(trade_data_by_sid) - } final_prices = { asset.sid: trade_data_by_sid[asset.sid]. - loc[asset.end_date.value].close + loc[asset.end_date].close for asset in assets } else: @@ -3527,27 +3197,29 @@ class TestEquityAutoClose(TestCase): ) -class TestOrderAfterDelist(TestCase): - @classmethod - def setUpClass(cls): - cls.env = TradingEnvironment() +class TestOrderAfterDelist(WithTradingEnvironment, ZiplineTestCase): + start = pd.Timestamp('2016-01-05', tz='utc') + day_1 = pd.Timestamp('2016-01-06', tz='utc') + day_4 = pd.Timestamp('2016-01-11', tz='utc') + end = pd.Timestamp('2016-01-15', tz='utc') - cls.days = cls.env.days_in_range( - start=pd.Timestamp("2016-01-05", tz='UTC'), - end=pd.Timestamp("2016-01-15", tz='UTC') + @classmethod + def make_equity_info(cls): + return pd.DataFrame.from_dict( + { + 1: { + 'start_date': cls.start, + 'end_date': cls.day_1, + 'auto_close_date': cls.day_4, + 'symbol': "ASSET1" + }, + }, + orient='index', ) - # asset goes from 1/5 to 1/6. - # asset liquidate date is 1/11. - cls.env.write_data(equities_data={ - 1: { - 'start_date': cls.days[0], - 'end_date': cls.days[1], - 'auto_close_date': cls.days[4], - 'symbol': "ASSET1" - } - }) - + @classmethod + def init_class_fixtures(cls): + super(TestOrderAfterDelist, cls).init_class_fixtures() cls.data_portal = FakeDataPortal(cls.env) def test_order_in_quiet_period(self): diff --git a/tests/test_api_shim.py b/tests/test_api_shim.py index b04fb09d..09957090 100644 --- a/tests/test_api_shim.py +++ b/tests/test_api_shim.py @@ -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 ) diff --git a/tests/test_assets.py b/tests/test_assets.py index 40cfc33e..70be6ab7 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -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) diff --git a/tests/test_bar_data.py b/tests/test_bar_data.py index 08ae9fdb..d1dc2c2d 100644 --- a/tests/test_bar_data.py +++ b/tests/test_bar_data.py @@ -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( diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 2b2960f4..ef9b1843 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -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 diff --git a/tests/test_blotter.py b/tests/test_blotter.py index 4a558d85..f1576b1b 100644 --- a/tests/test_blotter.py +++ b/tests/test_blotter.py @@ -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), diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py index 6bb6e351..f12a0dc4 100644 --- a/tests/test_exception_handling.py +++ b/tests/test_exception_handling.py @@ -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', diff --git a/tests/test_execution_styles.py b/tests/test_execution_styles.py index 42696a7e..05d81698 100644 --- a/tests/test_execution_styles.py +++ b/tests/test_execution_styles.py @@ -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): """ diff --git a/tests/test_fetcher.py b/tests/test_fetcher.py index 228b7c8e..71f77965 100644 --- a/tests/test_fetcher.py +++ b/tests/test_fetcher.py @@ -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'), diff --git a/tests/test_finance.py b/tests/test_finance.py index 5ba73a96..0f20f015 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -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 diff --git a/tests/test_history.py b/tests/test_history.py index 9343f6c0..0aec4913 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -1,5 +1,4 @@ from textwrap import dedent -from unittest import TestCase from numbers import Real @@ -9,49 +8,47 @@ from numpy import nan from numpy.testing import assert_almost_equal from nose_parameterized import parameterized -from testfixtures import TempDirectory from zipline import TradingAlgorithm from zipline._protocol import handle_non_market_minutes from zipline.assets import Asset -from zipline.data.data_portal import DataPortal, DailyHistoryAggregator -from zipline.data.minute_bars import ( - BcolzMinuteBarReader, - BcolzMinuteBarWriter, - US_EQUITIES_MINUTES_PER_DAY -) -from zipline.data.us_equity_pricing import ( - SQLiteAdjustmentWriter, - SQLiteAdjustmentReader, - BcolzDailyBarReader) +from zipline.data.data_portal import DailyHistoryAggregator from zipline.errors import ( HistoryInInitialize, HistoryWindowStartsBeforeData, ) -from zipline.finance.trading import ( - TradingEnvironment, - SimulationParameters -) +from zipline.finance.trading import SimulationParameters from zipline.protocol import BarData -from zipline.testing import str_to_seconds -from zipline.testing.core import ( - write_minute_data_for_asset, - DailyBarWriterFromDataFrames, - MockDailyBarReader +from zipline.testing import ( + create_minute_df_for_asset, + str_to_seconds, + MockDailyBarReader, ) from zipline.testing.fixtures import ( - WithBcolzMinutes, + WithBcolzMinuteBarReader, + WithDataPortal, ZiplineTestCase ) -OHLC = ["open", "high", "low", "close"] -OHLCV = OHLC + ["volume"] -OHLCP = OHLC + ["price"] -ALL_FIELDS = OHLCP + ["volume"] +OHLC = ['open', 'high', 'low', 'close'] +OHLCV = OHLC + ['volume'] +OHLCP = OHLC + ['price'] +ALL_FIELDS = OHLCP + ['volume'] -class HistoryTestCaseBase(TestCase): +class WithHistory(WithDataPortal): + TRADING_START_DT = TRADING_ENV_MIN_DATE = START_DATE = pd.Timestamp( + '2014-02-03', + tz='UTC', + ) + TRADING_END_DT = END_DATE = pd.Timestamp('2016-01-29', tz='UTC') + + SPLIT_ASSET_SID = 4 + DIVIDEND_ASSET_SID = 5 + MERGER_ASSET_SID = 6 + HALF_DAY_TEST_ASSET_SID = 7 + SHORT_ASSET_SID = 8 # asset1: # - 2014-03-01 (rounds up to TRADING_START_DT) to 2016-01-29. # - every minute/day. @@ -79,185 +76,160 @@ class HistoryTestCaseBase(TestCase): # - trades every minute # - merger on 2015-01-05 and 2015-01-06 @classmethod - def setUpClass(cls): - cls.tempdir = TempDirectory() - - # trading_start (when bcolz files start) = 2014-02-01 - cls.TRADING_START_DT = pd.Timestamp("2014-02-03", tz='UTC') - cls.TRADING_END_DT = pd.Timestamp("2016-01-29", tz='UTC') - - cls.env = TradingEnvironment(min_date=cls.TRADING_START_DT) - + def init_class_fixtures(cls): + super(WithHistory, cls).init_class_fixtures() cls.trading_days = cls.env.days_in_range( start=cls.TRADING_START_DT, end=cls.TRADING_END_DT ) - cls.create_assets() - - 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) - cls.SPLIT_ASSET = cls.env.asset_finder.retrieve_asset(4) - cls.DIVIDEND_ASSET = cls.env.asset_finder.retrieve_asset(5) - cls.MERGER_ASSET = cls.env.asset_finder.retrieve_asset(6) - cls.HALF_DAY_TEST_ASSET = cls.env.asset_finder.retrieve_asset(7) - cls.SHORT_ASSET = cls.env.asset_finder.retrieve_asset(8) - - cls.adj_reader = cls.create_adjustments_reader() - - cls.create_data() - - @classmethod - def tearDownClass(cls): - del cls.adj_reader - cls.tempdir.cleanup() - - def setUp(self): - self.create_data_portal() - - def tearDown(self): - del self.data_portal - - def create_data_portal(self): - raise NotImplementedError() - - @classmethod - def create_data(cls): - raise NotImplementedError() - - @classmethod - def create_assets(cls): - jan_5_2015 = pd.Timestamp("2015-01-05", tz='UTC') - day_after_12312015 = cls.env.next_trading_day( - pd.Timestamp("2015-12-31", tz='UTC') + cls.ASSET1 = cls.asset_finder.retrieve_asset(1) + cls.ASSET2 = cls.asset_finder.retrieve_asset(2) + cls.ASSET3 = cls.asset_finder.retrieve_asset(3) + cls.SPLIT_ASSET = cls.asset_finder.retrieve_asset( + cls.SPLIT_ASSET_SID, + ) + cls.DIVIDEND_ASSET = cls.asset_finder.retrieve_asset( + cls.DIVIDEND_ASSET_SID, + ) + cls.MERGER_ASSET = cls.asset_finder.retrieve_asset( + cls.MERGER_ASSET_SID, + ) + cls.HALF_DAY_TEST_ASSET = cls.asset_finder.retrieve_asset( + cls.HALF_DAY_TEST_ASSET_SID, + ) + cls.SHORT_ASSET = cls.asset_finder.retrieve_asset( + cls.SHORT_ASSET_SID, ) - cls.env.write_data(equities_data={ - 1: { - "start_date": pd.Timestamp("2014-01-03", tz='UTC'), - "end_date": cls.TRADING_END_DT, - "symbol": "ASSET1" - }, - 2: { - "start_date": jan_5_2015, - "end_date": day_after_12312015, - "symbol": "ASSET2" - }, - 3: { - "start_date": jan_5_2015, - "end_date": day_after_12312015, - "symbol": "ASSET3" - }, - 4: { - "start_date": jan_5_2015, - "end_date": day_after_12312015, - "symbol": "SPLIT_ASSET" - }, - 5: { - "start_date": jan_5_2015, - "end_date": day_after_12312015, - "symbol": "DIVIDEND_ASSET" - }, - 6: { - "start_date": jan_5_2015, - "end_date": day_after_12312015, - "symbol": "MERGER_ASSET" - }, - 7: { - "start_date": pd.Timestamp("2014-07-02", tz='UTC'), - "end_date": day_after_12312015, - "symbol": "HALF_DAY_TEST_ASSET" - }, - 8: { - "start_date": pd.Timestamp("2015-01-05", tz='UTC'), - "end_date": pd.Timestamp("2015-01-06", tz='UTC'), - "symbol": "SHORT_ASSET" - } - }) - @classmethod - def create_adjustments_reader(cls): - path = cls.tempdir.getpath("test_adjustments.db") + def make_equity_info(cls): + jan_5_2015 = pd.Timestamp('2015-01-05', tz='UTC') + day_after_12312015 = pd.Timestamp('2016-01-04', tz='UTC') - adj_writer = SQLiteAdjustmentWriter( - path, - cls.env.trading_days, - MockDailyBarReader() + return pd.DataFrame.from_dict( + { + 1: { + 'start_date': pd.Timestamp('2014-01-03', tz='UTC'), + 'end_date': cls.TRADING_END_DT, + 'symbol': 'ASSET1' + }, + 2: { + 'start_date': jan_5_2015, + 'end_date': day_after_12312015, + 'symbol': 'ASSET2' + }, + 3: { + 'start_date': jan_5_2015, + 'end_date': day_after_12312015, + 'symbol': 'ASSET3' + }, + cls.SPLIT_ASSET_SID: { + 'start_date': jan_5_2015, + 'end_date': day_after_12312015, + 'symbol': 'SPLIT_ASSET' + }, + cls.DIVIDEND_ASSET_SID: { + 'start_date': jan_5_2015, + 'end_date': day_after_12312015, + 'symbol': 'DIVIDEND_ASSET' + }, + cls.MERGER_ASSET_SID: { + 'start_date': jan_5_2015, + 'end_date': day_after_12312015, + 'symbol': 'MERGER_ASSET' + }, + cls.HALF_DAY_TEST_ASSET_SID: { + 'start_date': pd.Timestamp('2014-07-02', tz='UTC'), + 'end_date': day_after_12312015, + 'symbol': 'HALF_DAY_TEST_ASSET' + }, + cls.SHORT_ASSET_SID: { + 'start_date': pd.Timestamp('2015-01-05', tz='UTC'), + 'end_date': pd.Timestamp('2015-01-06', tz='UTC'), + 'symbol': 'SHORT_ASSET' + } + }, + orient='index', ) - splits = pd.DataFrame([ + @classmethod + def make_splits_data(cls): + return pd.DataFrame([ { - 'effective_date': str_to_seconds("2015-01-06"), + 'effective_date': str_to_seconds('2015-01-06'), 'ratio': 0.5, - 'sid': cls.SPLIT_ASSET.sid + 'sid': cls.SPLIT_ASSET_SID, }, { - 'effective_date': str_to_seconds("2015-01-07"), + 'effective_date': str_to_seconds('2015-01-07'), 'ratio': 0.5, - 'sid': cls.SPLIT_ASSET.sid + 'sid': cls.SPLIT_ASSET_SID, }, ]) - mergers = pd.DataFrame([ + @classmethod + def make_mergers_data(cls): + return pd.DataFrame([ { - 'effective_date': str_to_seconds("2015-01-06"), + 'effective_date': str_to_seconds('2015-01-06'), 'ratio': 0.5, - 'sid': cls.MERGER_ASSET.sid + 'sid': cls.MERGER_ASSET_SID, }, { - 'effective_date': str_to_seconds("2015-01-07"), + 'effective_date': str_to_seconds('2015-01-07'), 'ratio': 0.5, - 'sid': cls.MERGER_ASSET.sid + 'sid': cls.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([ { # only care about ex date, the other dates don't matter here 'ex_date': - pd.Timestamp("2015-01-06", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-06', tz='UTC').to_datetime64(), 'record_date': - pd.Timestamp("2015-01-06", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-06', tz='UTC').to_datetime64(), 'declared_date': - pd.Timestamp("2015-01-06", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-06', tz='UTC').to_datetime64(), 'pay_date': - pd.Timestamp("2015-01-06", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-06', tz='UTC').to_datetime64(), 'amount': 2.0, - 'sid': cls.DIVIDEND_ASSET.sid + 'sid': cls.DIVIDEND_ASSET_SID, }, { 'ex_date': - pd.Timestamp("2015-01-07", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-07', tz='UTC').to_datetime64(), 'record_date': - pd.Timestamp("2015-01-07", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-07', tz='UTC').to_datetime64(), 'declared_date': - pd.Timestamp("2015-01-07", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-07', tz='UTC').to_datetime64(), 'pay_date': - pd.Timestamp("2015-01-07", tz='UTC').to_datetime64(), + pd.Timestamp('2015-01-07', tz='UTC').to_datetime64(), 'amount': 4.0, - 'sid': cls.DIVIDEND_ASSET.sid + 'sid': cls.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() def verify_regular_dt(self, idx, dt, mode, fields=None, assets=None): - if mode == "daily": - freq = "1d" + if mode == 'daily': + freq = '1d' else: - freq = "1m" + freq = '1m' fields = fields if fields is not None else ALL_FIELDS assets = assets if assets is not None else [self.ASSET2, self.ASSET3] @@ -298,7 +270,7 @@ class HistoryTestCaseBase(TestCase): np.full(10, np.nan), asset_series ) - elif field == "volume": + elif field == 'volume': if asset == self.ASSET2: # asset2 should have some zeros (instead of nans) np.testing.assert_array_equal( @@ -351,7 +323,7 @@ class HistoryTestCaseBase(TestCase): asset3_answer_key, asset_series ) - elif field == "volume": + elif field == 'volume': asset3_answer_key = np.zeros(10) asset3_answer_key[-position_from_end] = \ value_for_asset3 * 100 @@ -369,7 +341,7 @@ class HistoryTestCaseBase(TestCase): asset3_answer_key, asset_series ) - elif field == "price": + elif field == 'price': # price is always forward filled # asset2 has prices every minute, so it's easy @@ -463,77 +435,45 @@ def check_internal_consistency(bar_data, assets, fields, bar_count, freq): # for example, the open is always 1 higher than the close, the high # is always 2 higher than the close, etc. MINUTE_FIELD_INFO = { - "open": 1, - "high": 2, - "low": -1, - "close": 0, - "price": 0, - "volume": 0, # unused, later we'll multiply by 100 + 'open': 1, + 'high': 2, + 'low': -1, + 'close': 0, + 'price': 0, + 'volume': 0, # unused, later we'll multiply by 100 } -class MinuteEquityHistoryTestCase(HistoryTestCaseBase): - def create_data_portal(self): - self.data_portal = DataPortal( - self.env, - equity_minute_reader=BcolzMinuteBarReader(self.tempdir.path), - adjustment_reader=self.adj_reader - ) - +class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): @classmethod - def create_data(cls): - 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] - - writer = BcolzMinuteBarWriter( - cls.trading_days[0], - cls.tempdir.path, - market_opens, - market_closes, - US_EQUITIES_MINUTES_PER_DAY - ) - - write_minute_data_for_asset( - cls.env, - writer, - pd.Timestamp("2014-01-03", tz='UTC'), - pd.Timestamp("2016-01-30", tz='UTC'), - 1, - start_val=2 - ) - - for sid in [2, 4, 5, 6, cls.SHORT_ASSET.sid]: - asset = cls.env.asset_finder.retrieve_asset(sid) - write_minute_data_for_asset( + def make_minute_bar_data(cls): + data = {} + sids = {2, 4, 5, 6, cls.SHORT_ASSET_SID, cls.HALF_DAY_TEST_ASSET_SID} + for sid in sids: + asset = cls.asset_finder.retrieve_asset(sid) + data[sid] = create_minute_df_for_asset( cls.env, - writer, asset.start_date, asset.end_date, - asset.sid, - start_val=2 + start_val=2, ) - write_minute_data_for_asset( + data[1] = create_minute_df_for_asset( cls.env, - writer, - cls.HALF_DAY_TEST_ASSET.start_date, - cls.HALF_DAY_TEST_ASSET.end_date, - cls.HALF_DAY_TEST_ASSET.sid, - start_val=2 + pd.Timestamp('2014-01-03', tz='utc'), + pd.Timestamp('2016-01-30', tz='utc'), + start_val=2, ) - asset3 = cls.env.asset_finder.retrieve_asset(3) - write_minute_data_for_asset( + asset3 = cls.asset_finder.retrieve_asset(3) + data[3] = create_minute_df_for_asset( cls.env, - writer, asset3.start_date, asset3.end_date, - asset3.sid, + start_val=2, interval=10, - start_val=2 ) + return data def test_history_in_initialize(self): algo_text = dedent( @@ -554,7 +494,7 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): sim_params = SimulationParameters( period_start=start, period_end=end, - capital_base=float("1.0e5"), + capital_base=float('1.0e5'), data_frequency='minute', emission_rate='daily', env=self.env, @@ -574,22 +514,22 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # since asset2 and asset3 both started trading on 1/5/2015, let's do # some history windows that are completely before that minutes = self.env.market_minutes_for_day( - self.env.previous_trading_day(pd.Timestamp("2015-01-05", tz='UTC')) + self.env.previous_trading_day(pd.Timestamp('2015-01-05', tz='UTC')) )[0:60] for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute") + bar_data = BarData(self.data_portal, lambda: minute, 'minute') check_internal_consistency( - bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, "1m" + bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, '1m' ) for field in ALL_FIELDS: # OHLCP should be NaN # Volume should be 0 - asset2_series = bar_data.history(self.ASSET2, field, 10, "1m") - asset3_series = bar_data.history(self.ASSET3, field, 10, "1m") + asset2_series = bar_data.history(self.ASSET2, field, 10, '1m') + asset3_series = bar_data.history(self.ASSET3, field, 10, '1m') - if field == "volume": + if field == 'volume': np.testing.assert_array_equal(np.zeros(10), asset2_series) np.testing.assert_array_equal(np.zeros(10), asset3_series) else: @@ -622,49 +562,49 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): asset = self.env.asset_finder.retrieve_asset(sid) minutes = self.env.market_minutes_for_day( - pd.Timestamp("2015-01-05", tz='UTC') + pd.Timestamp('2015-01-05', tz='UTC') )[0:60] for idx, minute in enumerate(minutes): - self.verify_regular_dt(idx, minute, "minute", + self.verify_regular_dt(idx, minute, 'minute', assets=[asset], fields=[field]) def test_minute_midnight(self): - midnight = pd.Timestamp("2015-01-06", tz='UTC') + midnight = pd.Timestamp('2015-01-06', tz='UTC') last_minute = self.env.previous_open_and_close(midnight)[1] midnight_bar_data = \ - BarData(self.data_portal, lambda: midnight, "minute") + BarData(self.data_portal, lambda: midnight, 'minute') yesterday_bar_data = \ - BarData(self.data_portal, lambda: last_minute, "minute") + BarData(self.data_portal, lambda: last_minute, 'minute') with handle_non_market_minutes(midnight_bar_data): for field in ALL_FIELDS: np.testing.assert_array_equal( - midnight_bar_data.history(self.ASSET2, field, 30, "1m"), - yesterday_bar_data.history(self.ASSET2, field, 30, "1m") + midnight_bar_data.history(self.ASSET2, field, 30, '1m'), + yesterday_bar_data.history(self.ASSET2, field, 30, '1m') ) def test_minute_after_asset_stopped(self): # SHORT_ASSET's last day was 2015-01-06 # get some history windows that straddle the end minutes = self.env.market_minutes_for_day( - pd.Timestamp("2015-01-07", tz='UTC') + pd.Timestamp('2015-01-07', tz='UTC') )[0:60] for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute") + bar_data = BarData(self.data_portal, lambda: minute, 'minute') check_internal_consistency( - bar_data, self.SHORT_ASSET, ALL_FIELDS, 30, "1m" + bar_data, self.SHORT_ASSET, ALL_FIELDS, 30, '1m' ) # Reset data portal because it has advanced past next test date. - self.create_data_portal() + data_portal = self.make_data_portal() # choose a window that contains the last minute of the asset - bar_data = BarData(self.data_portal, lambda: minutes[15], "minute") + bar_data = BarData(data_portal, lambda: minutes[15], 'minute') # close high low open price volume # 2015-01-06 20:47:00+00:00 768 770 767 769 768 76800 @@ -698,18 +638,18 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # 2015-01-07 14:45:00+00:00 NaN NaN NaN NaN NaN 0 # 2015-01-07 14:46:00+00:00 NaN NaN NaN NaN NaN 0 - window = bar_data.history(self.SHORT_ASSET, ALL_FIELDS, 30, "1m") + window = bar_data.history(self.SHORT_ASSET, ALL_FIELDS, 30, '1m') # there should be 14 values and 16 NaNs/0s for field in ALL_FIELDS: - if field == "volume": + if field == 'volume': np.testing.assert_array_equal( range(76800, 78101, 100), - window["volume"][0:14] + window['volume'][0:14] ) np.testing.assert_array_equal( np.zeros(16), - window["volume"][-16:] + window['volume'][-16:] ) else: np.testing.assert_array_equal( @@ -723,11 +663,11 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # now do a smaller window that is entirely contained after the asset # ends - window = bar_data.history(self.SHORT_ASSET, ALL_FIELDS, 5, "1m") + window = bar_data.history(self.SHORT_ASSET, ALL_FIELDS, 5, '1m') for field in ALL_FIELDS: - if field == "volume": - np.testing.assert_array_equal(np.zeros(5), window["volume"]) + if field == 'volume': + np.testing.assert_array_equal(np.zeros(5), window['volume']) else: np.testing.assert_array_equal(np.full(5, np.nan), window[field]) @@ -736,7 +676,7 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # self.SPLIT_ASSET and self.MERGER_ASSET had splits/mergers # on 1/6 and 1/7 - jan5 = pd.Timestamp("2015-01-05", tz='UTC') + jan5 = pd.Timestamp('2015-01-05', tz='UTC') # the assets' close column starts at 2 on the first minute of # 1/5, then goes up one per minute forever @@ -747,8 +687,8 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): [asset], self.env.get_open_and_close(jan5)[1], 10, - "1m", - "close" + '1m', + 'close' )[asset] np.testing.assert_array_equal(np.array(range(382, 392)), window1) @@ -756,10 +696,10 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # straddling the first event window2 = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-06 14:35", tz='UTC'), + pd.Timestamp('2015-01-06 14:35', tz='UTC'), 10, - "1m", - "close" + '1m', + 'close' )[asset] # five minutes from 1/5 should be halved @@ -771,10 +711,10 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # straddling both events! window3 = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-07 14:35", tz='UTC'), + pd.Timestamp('2015-01-07 14:35', tz='UTC'), 400, # 5 minutes of 1/7, 390 of 1/6, and 5 minutes of 1/5 - "1m", - "close" + '1m', + 'close' )[asset] # first five minutes should be 387-391, but quartered @@ -785,7 +725,7 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # next 390 minutes should be 392-781, but halved np.testing.assert_array_equal( - np.array(range(392, 782), dtype="float64") / 2, + np.array(range(392, 782), dtype='float64') / 2, window3[5:395] ) @@ -795,10 +735,10 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # after last event window4 = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-07 14:40", tz='UTC'), + pd.Timestamp('2015-01-07 14:40', tz='UTC'), 5, - "1m", - "close" + '1m', + 'close' )[asset] # should not be adjusted, should be 787 to 791 @@ -810,10 +750,10 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # before any of the dividends window1 = self.data_portal.get_history_window( [self.DIVIDEND_ASSET], - pd.Timestamp("2015-01-05 21:00", tz='UTC'), + pd.Timestamp('2015-01-05 21:00', tz='UTC'), 10, - "1m", - "close" + '1m', + 'close' )[self.DIVIDEND_ASSET] np.testing.assert_array_equal(np.array(range(382, 392)), window1) @@ -821,16 +761,16 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # straddling the first dividend window2 = self.data_portal.get_history_window( [self.DIVIDEND_ASSET], - pd.Timestamp("2015-01-06 14:35", tz='UTC'), + pd.Timestamp('2015-01-06 14:35', tz='UTC'), 10, - "1m", - "close" + '1m', + 'close' )[self.DIVIDEND_ASSET] # first dividend is 2%, so the first five values should be 2% lower # than before np.testing.assert_array_almost_equal( - np.array(range(387, 392), dtype="float64") * 0.98, + np.array(range(387, 392), dtype='float64') * 0.98, window2[0:5] ) @@ -840,21 +780,21 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # straddling both dividends window3 = self.data_portal.get_history_window( [self.DIVIDEND_ASSET], - pd.Timestamp("2015-01-07 14:35", tz='UTC'), + pd.Timestamp('2015-01-07 14:35', tz='UTC'), 400, # 5 minutes of 1/7, 390 of 1/6, and 5 minutes of 1/5 - "1m", - "close" + '1m', + 'close' )[self.DIVIDEND_ASSET] # first five minute from 1/7 should be hit by 0.9408 (= 0.98 * 0.96) np.testing.assert_array_almost_equal( - np.around(np.array(range(387, 392), dtype="float64") * 0.9408, 3), + np.around(np.array(range(387, 392), dtype='float64') * 0.9408, 3), window3[0:5] ) # next 390 minutes should be hit by 0.96 (second dividend) np.testing.assert_array_almost_equal( - np.array(range(392, 782), dtype="float64") * 0.96, + np.array(range(392, 782), dtype='float64') * 0.96, window3[5:395] ) @@ -880,8 +820,8 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): def test_overnight_adjustments(self): # Should incorporate adjustments on midnight 01/06 - current_dt = pd.Timestamp("2015-01-06 8:45", tz='US/Eastern') - bar_data = BarData(self.data_portal, lambda: current_dt, "minute") + current_dt = pd.Timestamp('2015-01-06 8:45', tz='US/Eastern') + bar_data = BarData(self.data_portal, lambda: current_dt, 'minute') expected = { 'open': np.arange(383, 393) / 2.0, @@ -943,14 +883,14 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): # # five minutes into the day after the early close, get 20 1m bars - dt = pd.Timestamp("2014-07-07 13:35:00", tz='UTC') + dt = pd.Timestamp('2014-07-07 13:35:00', tz='UTC') window = self.data_portal.get_history_window( [self.HALF_DAY_TEST_ASSET], dt, 20, - "1m", - "close" + '1m', + 'close' )[self.HALF_DAY_TEST_ASSET] # 390 minutes for 7/2, 210 minutes for 7/3, 7/4-7/6 closed @@ -961,12 +901,12 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): self.assertEqual( window.index[-6], - pd.Timestamp("2014-07-03 17:00", tz='UTC') + pd.Timestamp('2014-07-03 17:00', tz='UTC') ) self.assertEqual( window.index[-5], - pd.Timestamp("2014-07-07 13:31", tz='UTC') + pd.Timestamp('2014-07-07 13:31', tz='UTC') ) def test_minute_different_lifetimes(self): @@ -986,8 +926,8 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): [self.ASSET1, self.ASSET2], self.env.get_open_and_close(day)[0], 100, - "1m", - "close" + '1m', + 'close' ) np.testing.assert_array_equal( @@ -1006,99 +946,65 @@ class MinuteEquityHistoryTestCase(HistoryTestCaseBase): self.TRADING_START_DT ) exp_msg = ( - "History window extends before 2014-02-03. To use this history " - "window, start the backtest on or after 2014-02-04." + 'History window extends before 2014-02-03. To use this history ' + 'window, start the backtest on or after 2014-02-04.' ) for field in OHLCP: with self.assertRaisesRegexp( HistoryWindowStartsBeforeData, exp_msg): self.data_portal.get_history_window( - [self.ASSET1], first_day_minutes[5], 15, "1m", "price" + [self.ASSET1], first_day_minutes[5], 15, '1m', 'price' )[self.ASSET1] -class DailyEquityHistoryTestCase(HistoryTestCaseBase): - def create_data_portal(self): - daily_path = self.tempdir.getpath("testdaily.bcolz") - - self.data_portal = DataPortal( - self.env, - equity_daily_reader=BcolzDailyBarReader(daily_path), - equity_minute_reader=BcolzMinuteBarReader(self.tempdir.path), - adjustment_reader=self.adj_reader +class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): + @classmethod + def make_daily_bar_data(cls): + yield 1, cls.create_df_for_asset( + cls.START_DATE, + pd.Timestamp('2016-01-30', tz='UTC') ) + yield 3, cls.create_df_for_asset( + pd.Timestamp('2015-01-05', tz='UTC'), + pd.Timestamp('2015-12-31', tz='UTC'), + interval=10, + force_zeroes=True + ) + yield cls.SHORT_ASSET_SID, cls.create_df_for_asset( + pd.Timestamp('2015-01-05', tz='UTC'), + pd.Timestamp('2015-01-06', tz='UTC'), + ) + + for sid in {2, 4, 5, 6}: + asset = cls.asset_finder.retrieve_asset(sid) + yield sid, cls.create_df_for_asset( + asset.start_date, + asset.end_date, + ) @classmethod - def create_data(cls): - path = cls.tempdir.getpath("testdaily.bcolz") - - dfs = { - 1: cls.create_df_for_asset( - cls.TRADING_START_DT, - pd.Timestamp("2016-01-30", tz='UTC') + def make_minute_bar_data(cls): + asset1 = cls.asset_finder.retrieve_asset(1) + asset2 = cls.asset_finder.retrieve_asset(2) + return { + asset1.sid: create_minute_df_for_asset( + cls.env, + asset1.start_date, + asset1.end_date, + start_val=2, ), - 3: cls.create_df_for_asset( - pd.Timestamp("2015-01-05", tz='UTC'), - pd.Timestamp("2015-12-31", tz='UTC'), - interval=10, - force_zeroes=True + asset2.sid: create_minute_df_for_asset( + cls.env, + asset2.start_date, + cls.env.previous_trading_day(asset2.end_date), + start_val=2, + minute_blacklist=[ + pd.Timestamp('2015-01-08 14:31', tz='UTC'), + pd.Timestamp('2015-01-08 21:00', tz='UTC'), + ], ), - cls.SHORT_ASSET.sid: cls.create_df_for_asset( - pd.Timestamp("2015-01-05", tz='UTC'), - pd.Timestamp("2015-01-06", tz='UTC'), - ) } - for sid in [2, 4, 5, 6]: - asset = cls.env.asset_finder.retrieve_asset(sid) - dfs[sid] = cls.create_df_for_asset( - asset.start_date, - asset.end_date - ) - - days = cls.env.days_in_range( - cls.TRADING_START_DT, - cls.TRADING_END_DT - ) - - daily_writer = DailyBarWriterFromDataFrames(dfs) - daily_writer.write(path, days, dfs) - - 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 - ) - - write_minute_data_for_asset( - cls.env, - minute_writer, - cls.ASSET1.start_date, - cls.ASSET1.end_date, - cls.ASSET1.sid, - start_val=2 - ) - - write_minute_data_for_asset( - cls.env, - minute_writer, - cls.ASSET2.start_date, - cls.ASSET2.end_date, - cls.ASSET2.sid, - start_val=2, - minute_blacklist=[ - pd.Timestamp('2015-01-08 14:31', tz='UTC'), - pd.Timestamp('2015-01-08 21:00', tz='UTC'), - ] - ) - @classmethod def create_df_for_asset(cls, start_day, end_day, interval=1, force_zeroes=False): @@ -1109,13 +1015,16 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # want to start with a 0 days_arr = np.array(range(2, days_count + 2)) - df = pd.DataFrame({ - "open": days_arr + 1, - "high": days_arr + 2, - "low": days_arr - 1, - "close": days_arr, - "volume": 100 * days_arr, - }) + df = pd.DataFrame( + { + 'open': days_arr + 1, + 'high': days_arr + 2, + 'low': days_arr - 1, + 'close': days_arr, + 'volume': 100 * days_arr, + }, + index=days, + ) if interval > 1: counter = 0 @@ -1123,31 +1032,29 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): df[counter:(counter + interval - 1)] = 0 counter += interval - df["day"] = [day.value for day in days] - return df def test_daily_before_assets_trading(self): # asset2 and asset3 both started trading in 2015 days = self.env.days_in_range( - start=pd.Timestamp("2014-12-15", tz='UTC'), - end=pd.Timestamp("2014-12-18", tz='UTC'), + start=pd.Timestamp('2014-12-15', tz='UTC'), + end=pd.Timestamp('2014-12-18', tz='UTC'), ) for idx, day in enumerate(days): - bar_data = BarData(self.data_portal, lambda: day, "daily") + bar_data = BarData(self.data_portal, lambda: day, 'daily') check_internal_consistency( - bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, "1d" + bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, '1d' ) for field in ALL_FIELDS: # OHLCP should be NaN # Volume should be 0 - asset2_series = bar_data.history(self.ASSET2, field, 10, "1d") - asset3_series = bar_data.history(self.ASSET3, field, 10, "1d") + asset2_series = bar_data.history(self.ASSET2, field, 10, '1d') + asset3_series = bar_data.history(self.ASSET3, field, 10, '1d') - if field == "volume": + if field == 'volume': np.testing.assert_array_equal(np.zeros(10), asset2_series) np.testing.assert_array_equal(np.zeros(10), asset3_series) else: @@ -1166,7 +1073,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # 10 days # get the first 30 days of 2015 - jan5 = pd.Timestamp("2015-01-04") + jan5 = pd.Timestamp('2015-01-04') days = self.env.days_in_range( start=jan5, @@ -1174,19 +1081,19 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): ) for idx, day in enumerate(days): - self.verify_regular_dt(idx, day, "daily") + self.verify_regular_dt(idx, day, 'daily') def test_daily_some_assets_stopped(self): # asset1 ends on 2016-01-30 # asset2 ends on 2015-12-13 bar_data = BarData(self.data_portal, - lambda: pd.Timestamp("2016-01-06", tz='UTC'), - "daily") + lambda: pd.Timestamp('2016-01-06', tz='UTC'), + 'daily') for field in OHLCP: window = bar_data.history( - [self.ASSET1, self.ASSET2], field, 15, "1d" + [self.ASSET1, self.ASSET2], field, 15, '1d' ) # last 2 values for asset2 should be NaN (# of days since asset2 @@ -1200,7 +1107,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): self.assertFalse(np.isnan(window[self.ASSET2][-3])) volume_window = bar_data.history( - [self.ASSET1, self.ASSET2], "volume", 15, "1d" + [self.ASSET1, self.ASSET2], 'volume', 15, '1d' ) np.testing.assert_array_equal( @@ -1214,20 +1121,20 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # SHORT_ASSET trades on 1/5, 1/6, that's it. days = self.env.days_in_range( - start=pd.Timestamp("2015-01-07", tz='UTC'), - end=pd.Timestamp("2015-01-08", tz='UTC') + start=pd.Timestamp('2015-01-07', tz='UTC'), + end=pd.Timestamp('2015-01-08', tz='UTC') ) # days has 1/7, 1/8 for idx, day in enumerate(days): - bar_data = BarData(self.data_portal, lambda: day, "daily") + bar_data = BarData(self.data_portal, lambda: day, 'daily') check_internal_consistency( - bar_data, self.SHORT_ASSET, ALL_FIELDS, 2, "1d" + bar_data, self.SHORT_ASSET, ALL_FIELDS, 2, '1d' ) for field in ALL_FIELDS: asset_series = bar_data.history( - self.SHORT_ASSET, field, 2, "1d" + self.SHORT_ASSET, field, 2, '1d' ) if idx == 0: @@ -1239,7 +1146,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): ) self.assertTrue(np.isnan(asset_series.iloc[1])) - elif field == "volume": + elif field == 'volume': self.assertEqual(300, asset_series.iloc[0]) self.assertEqual(0, asset_series.iloc[1]) else: @@ -1247,7 +1154,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): if field in OHLCP: self.assertTrue(np.isnan(asset_series.iloc[0])) self.assertTrue(np.isnan(asset_series.iloc[1])) - elif field == "volume": + elif field == 'volume': self.assertEqual(0, asset_series.iloc[0]) self.assertEqual(0, asset_series.iloc[1]) @@ -1259,20 +1166,20 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # before any of the adjustments window1 = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-05", tz='UTC'), + pd.Timestamp('2015-01-05', tz='UTC'), 1, - "1d", - "close" + '1d', + 'close' )[asset] np.testing.assert_array_equal(window1, [2]) window1_volume = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-05", tz='UTC'), + pd.Timestamp('2015-01-05', tz='UTC'), 1, - "1d", - "volume" + '1d', + 'volume' )[asset] np.testing.assert_array_equal(window1_volume, [200]) @@ -1280,10 +1187,10 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # straddling the first event window2 = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-06", tz='UTC'), + pd.Timestamp('2015-01-06', tz='UTC'), 2, - "1d", - "close" + '1d', + 'close' )[asset] # first value should be halved, second value unadjusted @@ -1291,10 +1198,10 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): window2_volume = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-06", tz='UTC'), + pd.Timestamp('2015-01-06', tz='UTC'), 2, - "1d", - "volume" + '1d', + 'volume' )[asset] if asset == self.SPLIT_ASSET: @@ -1306,20 +1213,20 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # straddling both events window3 = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-07", tz='UTC'), + pd.Timestamp('2015-01-07', tz='UTC'), 3, - "1d", - "close" + '1d', + 'close' )[asset] np.testing.assert_array_equal([0.5, 1.5, 4], window3) window3_volume = self.data_portal.get_history_window( [asset], - pd.Timestamp("2015-01-07", tz='UTC'), + pd.Timestamp('2015-01-07', tz='UTC'), 3, - "1d", - "volume" + '1d', + 'volume' )[asset] if asset == self.SPLIT_ASSET: @@ -1333,10 +1240,10 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # before any dividend window1 = self.data_portal.get_history_window( [self.DIVIDEND_ASSET], - pd.Timestamp("2015-01-05", tz='UTC'), + pd.Timestamp('2015-01-05', tz='UTC'), 1, - "1d", - "close" + '1d', + 'close' )[self.DIVIDEND_ASSET] np.testing.assert_array_equal(window1, [2]) @@ -1344,10 +1251,10 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # straddling the first dividend window2 = self.data_portal.get_history_window( [self.DIVIDEND_ASSET], - pd.Timestamp("2015-01-06", tz='UTC'), + pd.Timestamp('2015-01-06', tz='UTC'), 2, - "1d", - "close" + '1d', + 'close' )[self.DIVIDEND_ASSET] # first dividend is 2%, so the first value should be 2% lower than @@ -1357,10 +1264,10 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # straddling both dividends window3 = self.data_portal.get_history_window( [self.DIVIDEND_ASSET], - pd.Timestamp("2015-01-07", tz='UTC'), + pd.Timestamp('2015-01-07', tz='UTC'), 3, - "1d", - "close" + '1d', + 'close' )[self.DIVIDEND_ASSET] # second dividend is 0.96 @@ -1373,12 +1280,12 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # asset2 ends on 2016-01-04 bar_data = BarData(self.data_portal, - lambda: pd.Timestamp("2016-01-06 16:00", tz='UTC'), - "daily") + lambda: pd.Timestamp('2016-01-06 16:00', tz='UTC'), + 'daily') for field in OHLCP: window = bar_data.history( - [self.ASSET1, self.ASSET2], field, 15, "1d" + [self.ASSET1, self.ASSET2], field, 15, '1d' ) # last 2 values for asset2 should be NaN @@ -1391,7 +1298,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): self.assertFalse(np.isnan(window[self.ASSET2][-3])) volume_window = bar_data.history( - [self.ASSET1, self.ASSET2], "volume", 15, "1d" + [self.ASSET1, self.ASSET2], 'volume', 15, '1d' ) np.testing.assert_array_equal( @@ -1406,7 +1313,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # last day # January 2015 has both daily and minute data for ASSET2 - day = pd.Timestamp("2015-01-07", tz='UTC') + day = pd.Timestamp('2015-01-07', tz='UTC') minutes = self.env.market_minutes_for_day(day) # minute data, baseline: @@ -1421,13 +1328,13 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): [self.ASSET2], minute, 3, - "1d", + '1d', field )[self.ASSET2] self.assertEqual(len(window), 3) - if field == "volume": + if field == 'volume': self.assertEqual(window[0], 200) self.assertEqual(window[1], 300) else: @@ -1436,19 +1343,19 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): last_val = -1 - if field == "open": + if field == 'open': last_val = 783 - elif field == "high": + elif field == 'high': # since we increase monotonically, it's just the last # value last_val = 784 + idx - elif field == "low": + elif field == 'low': # since we increase monotonically, the low is the first # value of the day last_val = 781 - elif field == "close" or field == "price": + elif field == 'close' or field == 'price': last_val = 782 + idx - elif field == "volume": + elif field == 'volume': # for volume, we sum up all the minutely volumes so far # today @@ -1462,7 +1369,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): # last day # January 2015 has both daily and minute data for ASSET2 - day = pd.Timestamp("2015-01-08", tz='UTC') + day = pd.Timestamp('2015-01-08', tz='UTC') minutes = self.env.market_minutes_for_day(day) # minute data, baseline: @@ -1476,13 +1383,13 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): [self.ASSET2], minute, 3, - "1d", + '1d', field )[self.ASSET2] self.assertEqual(len(window), 3) - if field == "volume": + if field == 'volume': self.assertEqual(window[0], 300) self.assertEqual(window[1], 400) else: @@ -1491,12 +1398,12 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): last_val = -1 - if field == "open": + if field == 'open': if idx == 0: last_val = np.nan else: last_val = 1174.0 - elif field == "high": + elif field == 'high': # since we increase monotonically, it's just the last # value if idx == 0: @@ -1505,28 +1412,28 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): last_val = 1562.0 else: last_val = 1174.0 + idx - elif field == "low": + elif field == 'low': # since we increase monotonically, the low is the first # value of the day if idx == 0: last_val = np.nan else: last_val = 1172.0 - elif field == "close": + elif field == 'close': if idx == 0: last_val = np.nan elif idx == 389: last_val = 1172.0 + 388 else: last_val = 1172.0 + idx - elif field == "price": + elif field == 'price': if idx == 0: last_val = 4 elif idx == 389: last_val = 1172.0 + 388 else: last_val = 1172.0 + idx - elif field == "volume": + elif field == 'volume': # for volume, we sum up all the minutely volumes so far # today if idx == 0: @@ -1539,7 +1446,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): np.array(range(1173, 1172 + idx + 1)) * 100) np.testing.assert_almost_equal(window[-1], last_val, - err_msg="field={0} minute={1}". + err_msg='field={0} minute={1}'. format(field, minute)) def test_history_window_before_first_trading_day(self): @@ -1549,8 +1456,8 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): second_day = self.env.next_trading_day(self.TRADING_START_DT) exp_msg = ( - "History window extends before 2014-02-03. To use this history " - "window, start the backtest on or after 2014-02-07." + 'History window extends before 2014-02-03. To use this history ' + 'window, start the backtest on or after 2014-02-07.' ) with self.assertRaisesRegexp(HistoryWindowStartsBeforeData, exp_msg): @@ -1558,8 +1465,8 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): [self.ASSET1], second_day, 4, - "1d", - "price" + '1d', + 'price' )[self.ASSET1] with self.assertRaisesRegexp(HistoryWindowStartsBeforeData, exp_msg): @@ -1567,8 +1474,8 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): [self.ASSET1], second_day, 4, - "1d", - "volume" + '1d', + 'volume' )[self.ASSET1] # Use a minute to force minute mode. @@ -1580,8 +1487,8 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): [self.ASSET2], first_minute, 4, - "1d", - "close" + '1d', + 'close' )[self.ASSET2] def test_history_window_different_order(self): @@ -1615,7 +1522,7 @@ class DailyEquityHistoryTestCase(HistoryTestCaseBase): window_2[self.ASSET2].values) -class MinuteToDailyAggregationTestCase(WithBcolzMinutes, +class MinuteToDailyAggregationTestCase(WithBcolzMinuteBarReader, ZiplineTestCase): # March 2016 @@ -1626,8 +1533,13 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, # 20 21 22 23 24 25 26 # 27 28 29 30 31 - TRADING_ENV_MIN_DATE = pd.Timestamp("2016-03-01", tz="UTC") - TRADING_ENV_MAX_DATE = pd.Timestamp("2016-03-31", tz="UTC") + TRADING_ENV_MIN_DATE = START_DATE = pd.Timestamp( + '2016-03-01', tz='UTC', + ) + TRADING_ENV_MAX_DATE = END_DATE = pd.Timestamp( + '2016-03-31', tz='UTC', + ) + ASSET_FINDER_EQUITY_SIDS = 1, 2 minutes = pd.date_range('2016-03-15 9:31', '2016-03-15 9:36', @@ -1635,23 +1547,7 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, tz='US/Eastern').tz_convert('UTC') @classmethod - def make_equities_info(cls): - return pd.DataFrame.from_dict({ - 1: { - "start_date": pd.Timestamp("2016-03-01", tz="UTC"), - "end_date": pd.Timestamp("2016-03-31", tz="UTC"), - "symbol": "EQUITY1", - }, - 2: { - "start_date": pd.Timestamp("2016-03-01", tz='UTC'), - "end_date": pd.Timestamp("2016-03-31", tz='UTC'), - "symbol": "EQUITY2" - }, - }, - orient='index') - - @classmethod - def make_bcolz_minute_bar_data(cls): + def make_minute_bar_data(cls): return { # sid data is created so that at least one high is lower than a # previous high, and the inverse for low @@ -1667,15 +1563,16 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, ), # sid 2 is included to provide data on different bars than sid 1, # as will as illiquidty mid-day - 2: pd.DataFrame({ - 'open': [201.50, nan, 204.50, nan, 200.50, 202.50], - 'high': [201.90, nan, 204.90, nan, 200.90, 202.90], - 'low': [201.10, nan, 204.10, nan, 200.10, 202.10], - 'close': [201.30, nan, 203.50, nan, 200.30, 202.30], - 'volume': [2001, 0, 2004, 0, 2000, 2002], - }, + 2: pd.DataFrame( + { + 'open': [201.50, nan, 204.50, nan, 200.50, 202.50], + 'high': [201.90, nan, 204.90, nan, 200.90, 202.90], + 'low': [201.10, nan, 204.10, nan, 200.10, 202.10], + 'close': [201.30, nan, 203.50, nan, 200.30, 202.30], + 'volume': [2001, 0, 2004, 0, 2000, 2002], + }, index=cls.minutes, - ) + ), } expected_values = { @@ -1755,9 +1652,9 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, repeat_results.append(value) assert_almost_equal(results, self.expected_values[asset][field], - err_msg="sid={0} field={1}".format(asset, field)) + err_msg='sid={0} field={1}'.format(asset, field)) assert_almost_equal(repeat_results, self.expected_values[asset][field], - err_msg="sid={0} field={1}".format(asset, field)) + err_msg='sid={0} field={1}'.format(asset, field)) @parameterized.expand([ ('open_sid_1', 'open', 1), @@ -1785,7 +1682,7 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, self.assertIsInstance(value, Real) assert_almost_equal(value, self.expected_values[sid][field][i], - err_msg="sid={0} field={1} dt={2}".format( + err_msg='sid={0} field={1} dt={2}'.format( sid, field, minute)) # Call a second time with the same dt, to prevent regression @@ -1797,7 +1694,7 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, self.assertIsInstance(value, Real) assert_almost_equal(value, self.expected_values[sid][field][i], - err_msg="sid={0} field={1} dt={2}".format( + err_msg='sid={0} field={1} dt={2}'.format( sid, field, minute)) @parameterized.expand(OHLCV) @@ -1831,11 +1728,11 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, for asset in assets: assert_almost_equal(results[asset], self.expected_values[asset][field], - err_msg="sid={0} field={1}".format( + err_msg='sid={0} field={1}'.format( asset, field)) assert_almost_equal(repeat_results[asset], self.expected_values[asset][field], - err_msg="sid={0} field={1}".format( + err_msg='sid={0} field={1}'.format( asset, field)) @parameterized.expand(OHLCV) @@ -1856,7 +1753,7 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, assert_almost_equal( value, self.expected_values[asset][field][i], - err_msg="sid={0} field={1} dt={2}".format( + err_msg='sid={0} field={1} dt={2}'.format( asset, field, minute)) # Call a second time with the same dt, to prevent regression @@ -1872,5 +1769,5 @@ class MinuteToDailyAggregationTestCase(WithBcolzMinutes, assert_almost_equal( value, self.expected_values[asset][field][i], - err_msg="sid={0} field={1} dt={2}".format( + err_msg='sid={0} field={1} dt={2}'.format( asset, field, minute)) diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 4174ea52..b77b9c0c 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -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): """ diff --git a/tests/test_security_list.py b/tests/test_security_list.py index 46f3feac..2b0223b1 100644 --- a/tests/test_security_list.py +++ b/tests/test_security_list.py @@ -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(): diff --git a/tests/utils/test_factory.py b/tests/utils/test_factory.py deleted file mode 100644 index 828c31a9..00000000 --- a/tests/utils/test_factory.py +++ /dev/null @@ -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 - ) diff --git a/tests/utils/test_final.py b/tests/utils/test_final.py index cb27baf0..22aff7e0 100644 --- a/tests/utils/test_final.py +++ b/tests/utils/test_final.py @@ -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') diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 9649fe21..dea5c74e 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -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 diff --git a/zipline/assets/__init__.py b/zipline/assets/__init__.py index 51b4358a..cdba9766 100644 --- a/zipline/assets/__init__.py +++ b/zipline/assets/__init__.py @@ -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', diff --git a/zipline/assets/asset_db_migrations.py b/zipline/assets/asset_db_migrations.py index 18ed61d2..0ea453b1 100644 --- a/zipline/assets/asset_db_migrations.py +++ b/zipline/assets/asset_db_migrations.py @@ -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'], + ) diff --git a/zipline/assets/asset_db_schema.py b/zipline/assets/asset_db_schema.py index bc79f5e7..3cff1668 100644 --- a/zipline/assets/asset_db_schema.py +++ b/zipline/assets/asset_db_schema.py @@ -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, diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 76b82293..b1eda14c 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -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, + ) diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index b6d70dc0..ce498336 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -101,6 +101,8 @@ class AssetFinder(object): PERSISTENT_TOKEN = "" 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: diff --git a/zipline/assets/synthetic.py b/zipline/assets/synthetic.py new file mode 100644 index 00000000..b071d3d0 --- /dev/null +++ b/zipline/assets/synthetic.py @@ -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, + ) diff --git a/zipline/data/_adjustments.pyx b/zipline/data/_adjustments.pyx index 5fded76b..c7900385 100644 --- a/zipline/data/_adjustments.pyx +++ b/zipline/data/_adjustments.pyx @@ -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()) diff --git a/zipline/data/minute_bars.py b/zipline/data/minute_bars.py index 3916bd83..dca46e10 100644 --- a/zipline/data/minute_bars.py +++ b/zipline/data/minute_bars.py @@ -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([ diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index ac1d3bb6..a9b41bb3 100644 --- a/zipline/data/us_equity_pricing.py +++ b/zipline/data/us_equity_pricing.py @@ -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. diff --git a/zipline/dispatch.py b/zipline/dispatch.py new file mode 100644 index 00000000..03c8d5f7 --- /dev/null +++ b/zipline/dispatch.py @@ -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 diff --git a/zipline/examples/buy_and_hold.py b/zipline/examples/buy_and_hold.py index f98738a8..69af2159 100644 --- a/zipline/examples/buy_and_hold.py +++ b/zipline/examples/buy_and_hold.py @@ -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) diff --git a/zipline/examples/buyapple.py b/zipline/examples/buyapple.py index 8984cd18..2ea1f0af 100755 --- a/zipline/examples/buyapple.py +++ b/zipline/examples/buyapple.py @@ -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) diff --git a/zipline/examples/dual_ema_talib.py b/zipline/examples/dual_ema_talib.py index 2a591bf6..92b1b41b 100755 --- a/zipline/examples/dual_ema_talib.py +++ b/zipline/examples/dual_ema_talib.py @@ -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. diff --git a/zipline/examples/dual_moving_average.py b/zipline/examples/dual_moving_average.py index af3d8609..e08c40c3 100755 --- a/zipline/examples/dual_moving_average.py +++ b/zipline/examples/dual_moving_average.py @@ -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. diff --git a/zipline/examples/olmar.py b/zipline/examples/olmar.py index b7199d34..ca116e13 100644 --- a/zipline/examples/olmar.py +++ b/zipline/examples/olmar.py @@ -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. diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 0837e441..e5978731 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -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 = "" - 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') diff --git a/zipline/pipeline/loaders/synthetic.py b/zipline/pipeline/loaders/synthetic.py index a212b2c0..913dac82 100644 --- a/zipline/pipeline/loaders/synthetic.py +++ b/zipline/pipeline/loaders/synthetic.py @@ -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): diff --git a/zipline/sources/__init__.py b/zipline/sources/__init__.py index e88807cb..343cce84 100644 --- a/zipline/sources/__init__.py +++ b/zipline/sources/__init__.py @@ -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', diff --git a/zipline/testing/__init__.py b/zipline/testing/__init__.py index 1946ec0d..5106cac5 100644 --- a/zipline/testing/__init__.py +++ b/zipline/testing/__init__.py @@ -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, ) diff --git a/zipline/testing/core.py b/zipline/testing/core.py index 8a99b8d2..f1109d75 100644 --- a/zipline/testing/core.py +++ b/zipline/testing/core.py @@ -1,53 +1,53 @@ +from abc import ABCMeta, abstractmethod, abstractproperty from contextlib import contextmanager from functools import wraps +import gzip from inspect import getargspec from itertools import ( combinations, count, product, ) -from nose.tools import nottest import operator import os +from os.path import abspath, dirname, join, realpath import shutil -from string import ascii_uppercase import tempfile -from bcolz import ctable -from logbook import FileHandler, TestHandler +from logbook import TestHandler from mock import patch +from nose.tools import nottest from numpy.testing import assert_allclose, assert_array_equal import pandas as pd -from pandas.tseries.offsets import MonthBegin -from six import iteritems, itervalues +from six import itervalues, iteritems, with_metaclass from six.moves import filter, map from sqlalchemy import create_engine +from testfixtures import TempDirectory from toolz import concat -from zipline.assets import AssetFinder -from zipline.assets.asset_writer import AssetDBWriterFromDataFrame -from zipline.assets.futures import CME_CODE_TO_MONTH +from zipline.assets import AssetFinder, AssetDBWriter +from zipline.assets.synthetic import make_simple_equity_info from zipline.data.data_portal import DataPortal from zipline.data.minute_bars import ( BcolzMinuteBarReader, BcolzMinuteBarWriter, US_EQUITIES_MINUTES_PER_DAY ) -from zipline.data.us_equity_pricing import SQLiteAdjustmentWriter, OHLC, \ - UINT32_MAX, BcolzDailyBarWriter, BcolzDailyBarReader +from zipline.data.us_equity_pricing import ( + BcolzDailyBarReader, + BcolzDailyBarWriter, + SQLiteAdjustmentWriter, +) +from zipline.finance.trading import TradingEnvironment from zipline.finance.order import ORDER_STATUS from zipline.pipeline.engine import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader -from zipline.finance.trading import TradingEnvironment from zipline.utils import security_list from zipline.utils.input_validation import expect_dimensions +from zipline.utils.sentinel import sentinel from zipline.utils.tradingcalendar import trading_days import numpy as np -from numpy import ( - float64, - uint32, - int64, -) +from numpy import float64 EPOCH = pd.Timestamp(0, tz='UTC') @@ -75,16 +75,6 @@ def str_to_seconds(s): return int((pd.Timestamp(s, tz='UTC') - EPOCH).total_seconds()) -def setup_logger(test, path='test.log'): - test.log_handler = FileHandler(path) - test.log_handler.push_application() - - -def teardown_logger(test): - test.log_handler.pop_application() - test.log_handler.close() - - def drain_zipline(test, zipline): output = [] transaction_count = 0 @@ -311,132 +301,6 @@ def chrange(start, stop): return list(map(chr, range(ord(start), ord(stop) + 1))) -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 - end_date : pd.Timestamp - 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': symbols, - 'start_date': [start_date] * num_assets, - 'end_date': [end_date] * num_assets, - 'exchange': 'TEST', - }, - index=sids, - ) - - -def make_jagged_equity_info(num_assets, - start_date, - first_end, - frequency, - periods_between_ends, - auto_close_delta): - """ - Create a dictionary 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`. - auto_close_delta : int - The constant delta that cascades the auto_close_date - Returns - ------- - info : pd.DataFrame - DataFrame representing newly-created assets. - """ - equity_info = {} - - for sid in range(num_assets): - equity_info[sid] = { - 'symbol': chr(ord('A') + sid), - 'start_date': start_date, - 'end_date': first_end + sid * periods_between_ends * frequency, - 'exchange': 'TEST', - 'auto_close_date': - auto_close_delta + first_end + sid * periods_between_ends * - frequency - } - - return equity_info - - def make_trade_data_for_asset_info(dates, asset_info, price_start, @@ -448,14 +312,12 @@ def make_trade_data_for_asset_info(dates, frequency, writer=None): """ - Convert the asset info dict to a dataframe of trade data for each sid, and - write to the writer if provided. Write NaNs for locations where assets did - not exist. Return a dict of the dataframes, keyed by sid. - locations where assets did not exist. + Convert the asset info dataframe into a dataframe of trade data for each + sid, and write to the writer if provided. Write NaNs for locations where + assets did not exist. Return a dict of the dataframes, keyed by sid. """ trade_data = {} - sids = asset_info.keys() - date_field = 'day' if frequency == 'daily' else 'dt' + sids = asset_info.index price_sid_deltas = np.arange(len(sids), dtype=float64) * price_step_by_sid price_date_deltas = (np.arange(len(dates), dtype=float64) * @@ -467,31 +329,24 @@ def make_trade_data_for_asset_info(dates, volumes = (volume_sid_deltas + volume_date_deltas[:, None]) + volume_start for j, sid in enumerate(sids): - start_date = asset_info[sid]['start_date'] - end_date = asset_info[sid]['end_date'] + start_date, end_date = asset_info.loc[sid, ['start_date', 'end_date']] # Normalize here so the we still generate non-NaN values on the minutes # for an asset's last trading day. for i, date in enumerate(dates.normalize()): if not (start_date <= date <= end_date): - prices[i, j] = np.nan + prices[i, j] = 0 volumes[i, j] = 0 - if frequency == 'daily': - dates_arr = [date.value for date in dates] - else: - dates_arr = dates - - df = pd.DataFrame({ - "open": prices[:, j], - "high": prices[:, j], - "low": prices[:, j], - "close": prices[:, j], - "volume": volumes[:, j], - date_field: dates_arr - }) - - if frequency == 'minute': - df = df.set_index(date_field) + df = pd.DataFrame( + { + "open": prices[:, j], + "high": prices[:, j], + "low": prices[:, j], + "close": prices[:, j], + "volume": volumes[:, j], + }, + index=dates, + ) if writer: writer.write(sid, df) @@ -501,148 +356,6 @@ def make_trade_data_for_asset_info(dates, return trade_data -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 - AssetDBWriterFromDataFrame. - """ - 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, - ) - - -def make_simple_asset_info(assets, start_date, end_date, symbols=None): - """ - Create a DataFrame representing assets that exist for the full duration - between `start_date` and `end_date`. - Parameters - ---------- - assets : array-like - start_date : pd.Timestamp - end_date : pd.Timestamp - 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(assets) - if symbols is None: - symbols = list(ascii_uppercase[:num_assets]) - return pd.DataFrame( - { - 'sid': assets, - 'symbol': symbols, - 'asset_type': ['equity'] * num_assets, - 'start_date': [start_date] * num_assets, - 'end_date': [end_date] * num_assets, - 'exchange': 'TEST', - } - ) - - def check_allclose(actual, desired, rtol=1e-07, @@ -699,87 +412,53 @@ class ExplodingObject(object): raise UnexpectedAttributeAccess(name) -class DailyBarWriterFromDataFrames(BcolzDailyBarWriter): - _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): - for asset in assets: - yield asset, ctable.fromdataframe(assets[asset]) - - 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) - ) - - def write_minute_data(env, tempdir, minutes, sids): - assets = {} - - length = len(minutes) - - for sid_idx, sid in enumerate(sids): - assets[sid] = pd.DataFrame({ - "open": (np.array(range(10, 10 + length)) + sid_idx), - "high": (np.array(range(15, 15 + length)) + sid_idx), - "low": (np.array(range(8, 8 + length)) + sid_idx), - "close": (np.array(range(10, 10 + length)) + sid_idx), - "volume": np.array(range(100, 100 + length)) + sid_idx, - "dt": minutes - }).set_index("dt") - write_bcolz_minute_data( env, env.days_in_range(minutes[0], minutes[-1]), tempdir.path, - assets + create_minute_bar_data(minutes, sids), ) - return tempdir.path +def create_minute_bar_data(minutes, sids): + length = len(minutes) + return { + sid: pd.DataFrame( + { + 'open': np.arange(length) + 10 + sid_idx, + 'high': np.arange(length) + 15 + sid_idx, + 'low': np.arange(length) + 8 + sid_idx, + 'close': np.arange(length) + 10 + sid_idx, + 'volume': np.arange(length) + 100 + sid_idx, + }, + index=minutes, + ) + for sid_idx, sid in enumerate(sids) + } + + +def create_daily_bar_data(trading_days, sids): + length = len(trading_days) + for sid_idx, sid in enumerate(sids): + yield sid, pd.DataFrame( + { + "open": (np.array(range(10, 10 + length)) + sid_idx), + "high": (np.array(range(15, 15 + length)) + sid_idx), + "low": (np.array(range(8, 8 + length)) + sid_idx), + "close": (np.array(range(10, 10 + length)) + sid_idx), + "volume": np.array(range(100, 100 + length)) + sid_idx, + "day": [day.value for day in trading_days] + }, + index=trading_days, + ) + + def write_daily_data(tempdir, sim_params, sids): path = os.path.join(tempdir.path, "testdaily.bcolz") - assets = {} - length = sim_params.days_in_period - for sid_idx, sid in enumerate(sids): - assets[sid] = pd.DataFrame({ - "open": (np.array(range(10, 10 + length)) + sid_idx), - "high": (np.array(range(15, 15 + length)) + sid_idx), - "low": (np.array(range(8, 8 + length)) + sid_idx), - "close": (np.array(range(10, 10 + length)) + sid_idx), - "volume": np.array(range(100, 100 + length)) + sid_idx, - "day": [day.value for day in sim_params.trading_days] - }, index=sim_params.trading_days) - - DailyBarWriterFromDataFrames(assets).write( - path, - sim_params.trading_days, - assets + BcolzDailyBarWriter(path, sim_params.trading_days).write( + create_daily_bar_data(sim_params.trading_days, sids), ) return path @@ -829,22 +508,27 @@ def write_bcolz_minute_data(env, days, path, df_dict): writer.write(sid, df) -def write_minute_data_for_asset(env, writer, start_dt, end_dt, sid, - interval=1, start_val=1, - minute_blacklist=None): +def create_minute_df_for_asset(env, + start_dt, + end_dt, + interval=1, + start_val=1, + minute_blacklist=None): asset_minutes = env.minutes_for_days_in_range(start_dt, end_dt) minutes_count = len(asset_minutes) minutes_arr = np.array(range(start_val, start_val + minutes_count)) - df = pd.DataFrame({ - "open": minutes_arr + 1, - "high": minutes_arr + 2, - "low": minutes_arr - 1, - "close": minutes_arr, - "volume": 100 * minutes_arr, - "dt": asset_minutes - }).set_index("dt") + df = pd.DataFrame( + { + "open": minutes_arr + 1, + "high": minutes_arr + 2, + "low": minutes_arr - 1, + "close": minutes_arr, + "volume": 100 * minutes_arr, + }, + index=asset_minutes, + ) if interval > 1: counter = 0 @@ -856,22 +540,24 @@ def write_minute_data_for_asset(env, writer, start_dt, end_dt, sid, for minute in minute_blacklist: df.loc[minute] = 0 - writer.write(sid, df) + return df def create_daily_df_for_asset(env, start_day, end_day, interval=1): days = env.days_in_range(start_day, end_day) days_count = len(days) - days_arr = np.array(range(2, days_count + 2)) + days_arr = np.arange(days_count) + 2 - df = pd.DataFrame({ - "open": days_arr + 1, - "high": days_arr + 2, - "low": days_arr - 1, - "close": days_arr, - "volume": days_arr * 100, - "day": [day.value for day in days] - }) + df = pd.DataFrame( + { + "open": days_arr + 1, + "high": days_arr + 2, + "low": days_arr - 1, + "close": days_arr, + "volume": days_arr * 100, + }, + index=days, + ) if interval > 1: # only keep every 'interval' rows @@ -886,37 +572,38 @@ def create_daily_df_for_asset(env, start_day, end_day, interval=1): return df +def trades_by_sid_to_dfs(trades_by_sid, index): + for sidint, trades in iteritems(trades_by_sid): + opens = [] + highs = [] + lows = [] + closes = [] + volumes = [] + for trade in trades: + opens.append(trade["open_price"]) + highs.append(trade["high"]) + lows.append(trade["low"]) + closes.append(trade["close_price"]) + volumes.append(trade["volume"]) + + yield sidint, pd.DataFrame( + { + "open": opens, + "high": highs, + "low": lows, + "close": closes, + "volume": volumes, + }, + index=index, + ) + + def create_data_portal_from_trade_history(env, tempdir, sim_params, trades_by_sid): if sim_params.data_frequency == "daily": path = os.path.join(tempdir.path, "testdaily.bcolz") - assets = {} - for sidint, trades in iteritems(trades_by_sid): - opens = [] - highs = [] - lows = [] - closes = [] - volumes = [] - for trade in trades: - opens.append(trade["open_price"]) - highs.append(trade["high"]) - lows.append(trade["low"]) - closes.append(trade["close_price"]) - volumes.append(trade["volume"]) - - assets[sidint] = pd.DataFrame({ - "open": np.array(opens), - "high": np.array(highs), - "low": np.array(lows), - "close": np.array(closes), - "volume": np.array(volumes), - "day": [day.value for day in sim_params.trading_days] - }, index=sim_params.trading_days) - - DailyBarWriterFromDataFrames(assets).write( - path, - sim_params.trading_days, - assets + BcolzDailyBarWriter(path, sim_params.trading_days).write( + trades_by_sid_to_dfs(trades_by_sid, sim_params.trading_days), ) equity_daily_reader = BcolzDailyBarReader(path) @@ -1039,30 +726,50 @@ class tmp_assets_db(object): Parameters ---------- - data : pd.DataFrame, optional - The data to feed to the writer. By default this maps: + **frames + The frames to pass to the AssetDBWriter. + By default this maps equities: ('A', 'B', 'C') -> map(ord, 'ABC') + + See Also + -------- + empty_assets_db + tmp_asset_finder """ - def __init__(self, **frames): + _default_equities = sentinel('_default_equities') + + def __init__(self, equities=_default_equities, **frames): self._eng = None - if not frames: - frames = { - 'equities': make_simple_equity_info( - list(map(ord, 'ABC')), - pd.Timestamp(0), - pd.Timestamp('2015'), - ) - } - self._data = AssetDBWriterFromDataFrame(**frames) + if equities is self._default_equities: + equities = make_simple_equity_info( + list(map(ord, 'ABC')), + pd.Timestamp(0), + pd.Timestamp('2015'), + ) + + frames['equities'] = equities + self._frames = frames + self._eng = None # set in enter and exit def __enter__(self): self._eng = eng = create_engine('sqlite://') - self._data.write_all(eng) + AssetDBWriter(eng).write(**self._frames) return eng def __exit__(self, *excinfo): assert self._eng is not None, '_eng was not set in __enter__' self._eng.dispose() + self._eng = None + + +def empty_assets_db(): + """Context manager for creating an empty assets db. + + See Also + -------- + tmp_assets_db + """ + return tmp_assets_db(equities=None) class tmp_asset_finder(tmp_assets_db): @@ -1070,8 +777,14 @@ class tmp_asset_finder(tmp_assets_db): Parameters ---------- - data : dict, optional - The data to feed to the writer + finder_cls : type, optional + The type of asset finder to create from the assets db. + **frames + Forwarded to ``tmp_assets_db``. + + See Also + -------- + tmp_assets_db """ def __init__(self, finder_cls=AssetFinder, **frames): self._finder_cls = finder_cls @@ -1081,6 +794,43 @@ class tmp_asset_finder(tmp_assets_db): return self._finder_cls(super(tmp_asset_finder, self).__enter__()) +def empty_asset_finder(): + """Context manager for creating an empty asset finder. + + See Also + -------- + empty_assets_db + tmp_assets_db + tmp_asset_finder + """ + return tmp_asset_finder(equities=None) + + +class tmp_trading_env(tmp_asset_finder): + """Create a temporary trading environment. + + Parameters + ---------- + finder_cls : type, optional + The type of asset finder to create from the assets db. + **frames + Forwarded to ``tmp_assets_db``. + + See Also + -------- + empty_trading_env + tmp_asset_finder + """ + def __enter__(self): + return TradingEnvironment( + asset_db_path=super(tmp_trading_env, self).__enter__().engine, + ) + + +def empty_trading_env(): + return tmp_trading_env(equities=None) + + class SubTestFailures(AssertionError): def __init__(self, *failures): self.failures = failures @@ -1178,48 +928,31 @@ class MockDailyBarReader(object): return 100 -def create_mock_adjustments(tempdir, days, splits=None, dividends=None, - mergers=None): - path = tempdir.getpath("test_adjustments.db") - - # create a split for the last day - writer = SQLiteAdjustmentWriter(path, days, MockDailyBarReader()) +def create_mock_adjustment_data(splits=None, dividends=None, mergers=None): if splits is None: splits = create_empty_splits_mergers_frame() - else: + elif not isinstance(splits, pd.DataFrame): splits = pd.DataFrame(splits) if mergers is None: mergers = create_empty_splits_mergers_frame() - else: + elif not isinstance(mergers, pd.DataFrame): mergers = pd.DataFrame(mergers) if dividends is None: - data = { - # Hackery to make the dtypes correct on an empty frame. - 'ex_date': np.array([], dtype='datetime64[ns]'), - 'pay_date': np.array([], dtype='datetime64[ns]'), - 'record_date': np.array([], dtype='datetime64[ns]'), - 'declared_date': np.array([], dtype='datetime64[ns]'), - 'amount': np.array([], dtype=float64), - 'sid': np.array([], dtype=int64), - } - dividends = pd.DataFrame( - data, - index=pd.DatetimeIndex([], tz='UTC'), - columns=['ex_date', - 'pay_date', - 'record_date', - 'declared_date', - 'amount', - 'sid'] - ) - else: - if not isinstance(dividends, pd.DataFrame): - dividends = pd.DataFrame(dividends) + dividends = create_empty_dividends_frame() + elif not isinstance(dividends, pd.DataFrame): + dividends = pd.DataFrame(dividends) - writer.write(splits, mergers, dividends) + return splits, mergers, dividends + +def create_mock_adjustments(tempdir, days, splits=None, dividends=None, + mergers=None): + path = tempdir.getpath("test_adjustments.db") + SQLiteAdjustmentWriter(path, MockDailyBarReader(), days).write( + *create_mock_adjustment_data(splits, dividends, mergers) + ) return path @@ -1359,16 +1092,34 @@ def parameter_space(**params): return decorator +def create_empty_dividends_frame(): + return pd.DataFrame( + np.array( + [], + dtype=[ + ('ex_date', 'datetime64[ns]'), + ('pay_date', 'datetime64[ns]'), + ('record_date', 'datetime64[ns]'), + ('declared_date', 'datetime64[ns]'), + ('amount', 'float64'), + ('sid', 'int32'), + ], + ), + index=pd.DatetimeIndex([], tz='UTC'), + ) + + def create_empty_splits_mergers_frame(): return pd.DataFrame( - { - # Hackery to make the dtypes correct on an empty frame. - 'effective_date': np.array([], dtype=int64), - 'ratio': np.array([], dtype=float64), - 'sid': np.array([], dtype=int64), - }, + np.array( + [], + dtype=[ + ('effective_date', 'int64'), + ('ratio', 'float64'), + ('sid', 'int64'), + ], + ), index=pd.DatetimeIndex([]), - columns=['effective_date', 'ratio', 'sid'], ) @@ -1409,3 +1160,156 @@ def make_test_handler(testcase, *args, **kwargs): handler = TestHandler(*args, **kwargs) testcase.addCleanup(handler.close) return handler + + +def write_compressed(path, content): + """ + Write a compressed (gzipped) file to `path`. + """ + with gzip.open(path, 'wb') as f: + f.write(content) + + +def read_compressed(path): + """ + Write a compressed (gzipped) file from `path`. + """ + with gzip.open(path, 'rb') as f: + return f.read() + + +zipline_git_root = abspath( + join(realpath(dirname(__file__)), '..', '..'), +) + + +def test_resource_path(*path_parts): + return os.path.join(zipline_git_root, 'tests', 'resources', *path_parts) + + +@contextmanager +def patch_os_environment(remove=None, **values): + """ + Context manager for patching the operating system environment. + """ + old_values = {} + remove = remove or [] + for key in remove: + old_values[key] = os.environ.pop(key) + for key, value in values.iteritems(): + old_values[key] = os.getenv(key) + os.environ[key] = value + try: + yield + finally: + for old_key, old_value in old_values.iteritems(): + if old_value is None: + # Value was not present when we entered, so del it out if it's + # still present. + try: + del os.environ[key] + except KeyError: + pass + else: + # Restore the old value. + os.environ[old_key] = old_value + + +class tmp_dir(TempDirectory, object): + """New style class that wrapper for TempDirectory in python 2. + """ + pass + + +class _TmpBarReader(with_metaclass(ABCMeta, tmp_dir)): + """A helper for tmp_bcolz_minute_bar_reader and tmp_bcolz_daily_bar_reader. + + Parameters + ---------- + env : TradingEnvironment + The trading env. + days : pd.DatetimeIndex + The days to write for. + data : dict[int -> pd.DataFrame] + The data to write. + path : str, optional + The path to the directory to write the data into. If not given, this + will be a unique name. + """ + @abstractproperty + def _reader_cls(self): + raise NotImplementedError('_reader') + + @abstractmethod + def _write(self, env, days, path, data): + raise NotImplementedError('_write') + + def __init__(self, env, days, data, path=None): + super(_TmpBarReader, self).__init__(path=path) + self._env = env + self._days = days + self._data = data + + def __enter__(self): + tmpdir = super(_TmpBarReader, self).__enter__() + env = self._env + try: + self._write( + env, + self._days, + tmpdir.path, + self._data, + ) + return self._reader_cls(tmpdir.path) + except: + self.__exit__(None, None, None) + raise + + +class tmp_bcolz_minute_bar_reader(_TmpBarReader): + """A temporary BcolzMinuteBarReader object. + + Parameters + ---------- + env : TradingEnvironment + The trading env. + days : pd.DatetimeIndex + The days to write for. + data : iterable[(int, pd.DataFrame)] + The data to write. + path : str, optional + The path to the directory to write the data into. If not given, this + will be a unique name. + + See Also + -------- + tmp_bcolz_daily_bar_reader + """ + _reader_cls = BcolzMinuteBarReader + _write = staticmethod(write_bcolz_minute_data) + + +class tmp_bcolz_daily_bar_reader(_TmpBarReader): + """A temporary BcolzDailyBarReader object. + + Parameters + ---------- + env : TradingEnvironment + The trading env. + days : pd.DatetimeIndex + The days to write for. + data : dict[int -> pd.DataFrame] + The data to write. + path : str, optional + The path to the directory to write the data into. If not given, this + will be a unique name. + + See Also + -------- + tmp_bcolz_daily_bar_reader + """ + _reader_cls = BcolzDailyBarReader + + @staticmethod + def _write(env, days, path, data): + BcolzDailyBarWriter(path, days).write(data) diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index 1d8d8d3d..1689dc8e 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -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(), + ) diff --git a/zipline/testing/predicates.py b/zipline/testing/predicates.py new file mode 100644 index 00000000..6dd6f8d9 --- /dev/null +++ b/zipline/testing/predicates.py @@ -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) diff --git a/zipline/utils/cache.py b/zipline/utils/cache.py index 9f9f1e53..657f17ed 100644 --- a/zipline/utils/cache.py +++ b/zipline/utils/cache.py @@ -1,5 +1,5 @@ """ -Cached object with an expiration date. +Caching utilities for zipline """ from collections import namedtuple diff --git a/zipline/utils/classproperty.py b/zipline/utils/classproperty.py new file mode 100644 index 00000000..59d1f367 --- /dev/null +++ b/zipline/utils/classproperty.py @@ -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) diff --git a/zipline/utils/cli.py b/zipline/utils/cli.py index 45f38499..4ff2138c 100644 --- a/zipline/utils/cli.py +++ b/zipline/utils/cli.py @@ -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) diff --git a/zipline/utils/compat.py b/zipline/utils/compat.py new file mode 100644 index 00000000..fff02e6b --- /dev/null +++ b/zipline/utils/compat.py @@ -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', +] diff --git a/zipline/utils/context_tricks.py b/zipline/utils/context_tricks.py index c4746616..a0a244e6 100644 --- a/zipline/utils/context_tricks.py +++ b/zipline/utils/context_tricks.py @@ -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): diff --git a/zipline/utils/control_flow.py b/zipline/utils/control_flow.py index 253a71b3..7f0ab83e 100644 --- a/zipline/utils/control_flow.py +++ b/zipline/utils/control_flow.py @@ -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)} diff --git a/zipline/utils/data.py b/zipline/utils/data.py index 9c10b8d7..f02e02b1 100644 --- a/zipline/utils/data.py +++ b/zipline/utils/data.py @@ -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(, - [(0, 'zero'), (-1, 'negative one'), (2, 'two')]) - >>> d[1] = 'one' # Mutating the mapping maintains the sort order. - >>> d # doctest: +NORMALIZE_WHITESPACE - SortedDict(, - [(0, 'zero'), (-1, 'negative one'), (1, 'one'), (2, 'two')]) - >>> del d[0] - >>> d # doctest: +NORMALIZE_WHITESPACE - SortedDict(, - [(-1, 'negative one'), (1, 'one'), (2, 'two')]) - >>> del d[2] - >>> d - SortedDict(, [(-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] diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index fe69133e..9adb4dcb 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -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'] diff --git a/zipline/utils/final.py b/zipline/utils/final.py index 02394535..2e262a9c 100644 --- a/zipline/utils/final.py +++ b/zipline/utils/final.py @@ -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)): diff --git a/zipline/utils/functional.py b/zipline/utils/functional.py index 05a1d632..79569de0 100644 --- a/zipline/utils/functional.py +++ b/zipline/utils/functional.py @@ -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 diff --git a/zipline/utils/input_validation.py b/zipline/utils/input_validation.py index 51708dd0..19a29f6e 100644 --- a/zipline/utils/input_validation.py +++ b/zipline/utils/input_validation.py @@ -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, + ), ) diff --git a/zipline/utils/metautils.py b/zipline/utils/metautils.py new file mode 100644 index 00000000..fa0dd0ed --- /dev/null +++ b/zipline/utils/metautils.py @@ -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, + {}, + ) diff --git a/zipline/utils/preprocess.py b/zipline/utils/preprocess.py index 04f070e3..29d74232 100644 --- a/zipline/utils/preprocess.py +++ b/zipline/utils/preprocess.py @@ -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( """\