From bc0b117dc94b8e93081818964e3b1bdbf9b33abb Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Thu, 31 Mar 2016 15:26:57 -0400 Subject: [PATCH] MAINT: make the data loading apis more consistent. Changes BcolzDailyBarWriter to not be an abc, data is passed as an iterator of (sid, dataframe) pairs to the write method. Changes the AssetsDBWriter to be a single class which accepts an engine at construction time and has a `write` method for writing dataframes for the various tables. We no longer support writing the various other data types, callers should coerce their data into a dataframe themselves. See zipline.assets.synthetic for some helpers to do this. Adds many new fixtures and updates some existing fixtures to use the new ones: WithDefaultDateBounds A fixture that provides the suite a START_DATE and END_DATE. This is meant to make it easy for other fixtures to synchronize their date ranges without depending on eachother in strange ways. For example, WithBcolzMinuteBarReader and WithBcolzDailyBarReader by default should both have data for the same dates, so they may use depend on WithDefaultDates without forcing a dependency between them. WithTmpDir, WithInstanceTmpDir Provides the suite or individual test case a temporary directory. WithBcolzDailyBarReader Provides the suite a BcolzDailyBarReader which reads from bcolz data written to a temporary directory. The data will be read from dataframes and then converted to bcolz files with BcolzDailyBarWriter.write WithBcolzDailyBarReaderFromCSVs Provides the suite a BcolzDailyBarReader which reads from bcolz data written to a temporary directory. The data will be read from a collection of CSV files and then converted into the bcolz data through BcolzDailyBarWriter.write_csvs WithBcolzMinuteBarReader Provides the suite a BcolzMinuteBarReader which reads from bcolz data written to a temporary directory. The data will be read from dataframes and then converted to bcolz files with BcolzMinuteBarWriter.write WithAdjustmentReader Provides the suite a SQLiteAdjustmentReader which reads from an in memory sqlite database. The data will be read from dataframes and then converted into sqlite with SQLiteAdjustmentWriter.write WithDataPortal Provides each test case a DataPortal object with data from temporary resources. --- .gitignore | 2 + docs/source/whatsnew/1.0.0.txt | 64 + etc/requirements.txt | 1 + tests/data/test_us_equity_pricing.py | 119 +- tests/finance/test_slippage.py | 208 +-- tests/pipeline/base.py | 4 +- tests/pipeline/test_blaze.py | 6 +- tests/pipeline/test_earnings.py | 1 - tests/pipeline/test_engine.py | 201 ++- tests/pipeline/test_pipeline_algo.py | 241 ++- .../pipeline/test_us_equity_pricing_loader.py | 112 +- .../fetcher_inputs/fetcher_test_data.py | 35 - tests/resources/quandl_samples/AAPL.csv.gz | Bin 0 -> 13510 bytes tests/resources/quandl_samples/BRK_A.csv.gz | Bin 0 -> 4799 bytes tests/resources/quandl_samples/MSFT.csv.gz | Bin 0 -> 11579 bytes tests/resources/quandl_samples/ZEN.csv.gz | Bin 0 -> 3007 bytes .../quandl_samples/rebuild_samples.py | 39 + tests/test_algorithm.py | 1302 ++++++----------- tests/test_api_shim.py | 194 +-- tests/test_assets.py | 534 +++---- tests/test_bar_data.py | 498 ++++--- tests/test_benchmark.py | 170 +-- tests/test_blotter.py | 115 +- tests/test_exception_handling.py | 46 +- tests/test_execution_styles.py | 19 +- tests/test_fetcher.py | 196 +-- tests/test_finance.py | 121 +- tests/test_history.py | 849 +++++------ tests/test_perf_tracking.py | 319 ++-- tests/test_security_list.py | 139 +- tests/utils/test_factory.py | 60 - tests/utils/test_final.py | 4 +- zipline/algorithm.py | 73 +- zipline/assets/__init__.py | 4 + zipline/assets/asset_db_migrations.py | 105 +- zipline/assets/asset_db_schema.py | 21 +- zipline/assets/asset_writer.py | 523 +++---- zipline/assets/assets.py | 73 +- zipline/assets/synthetic.py | 257 ++++ zipline/data/_adjustments.pyx | 6 + zipline/data/minute_bars.py | 27 +- zipline/data/us_equity_pricing.py | 640 ++++---- zipline/dispatch.py | 21 + zipline/examples/buy_and_hold.py | 3 +- zipline/examples/buyapple.py | 3 +- zipline/examples/dual_ema_talib.py | 3 +- zipline/examples/dual_moving_average.py | 3 +- zipline/examples/olmar.py | 4 +- zipline/finance/trading.py | 194 +-- zipline/pipeline/loaders/synthetic.py | 178 ++- zipline/sources/__init__.py | 7 +- zipline/testing/__init__.py | 30 +- zipline/testing/core.py | 886 +++++------ zipline/testing/fixtures.py | 740 ++++++++-- zipline/testing/predicates.py | 140 ++ zipline/utils/cache.py | 2 +- zipline/utils/classproperty.py | 8 + zipline/utils/cli.py | 32 + zipline/utils/compat.py | 19 + zipline/utils/context_tricks.py | 8 + zipline/utils/control_flow.py | 10 + zipline/utils/data.py | 88 -- zipline/utils/factory.py | 8 +- zipline/utils/final.py | 67 +- zipline/utils/functional.py | 161 +- zipline/utils/input_validation.py | 5 +- zipline/utils/metautils.py | 74 + zipline/utils/preprocess.py | 45 +- 68 files changed, 5145 insertions(+), 4922 deletions(-) create mode 100644 docs/source/whatsnew/1.0.0.txt create mode 100644 tests/resources/quandl_samples/AAPL.csv.gz create mode 100644 tests/resources/quandl_samples/BRK_A.csv.gz create mode 100644 tests/resources/quandl_samples/MSFT.csv.gz create mode 100644 tests/resources/quandl_samples/ZEN.csv.gz create mode 100644 tests/resources/quandl_samples/rebuild_samples.py delete mode 100644 tests/utils/test_factory.py create mode 100644 zipline/assets/synthetic.py create mode 100644 zipline/dispatch.py create mode 100644 zipline/testing/predicates.py create mode 100644 zipline/utils/classproperty.py create mode 100644 zipline/utils/compat.py create mode 100644 zipline/utils/metautils.py 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 0000000000000000000000000000000000000000..4c2c4e23bc39592c65c17248c836091685a0b825 GIT binary patch literal 13510 zcmV;%G&##3iwFpMfpAs=|3N`eOfF+{b^wK)+p26$a-Hw>6ncQ_MV=3r#xnTAV6egW z0yTCESuL=9czlm<#H^JyXZ;}P{@0$n*3OI>8HW*vjQYF3|JOgi{`0^5`JZ3^@h|`V zf4u(FU;ppx-~ZEJ|JR>i|Lw2;^uPb-pI`ss|NZIj{_=nR@(+Lh=YM$p*MIq^zx?ap z{+GZ1*T4Mr_3!@SAAkPV|IUB%KlYzk!GDNl{Kr`9Z%WGTPbvQ?l~?cmXDt1gKdI_R z`AH?~XZ^|bwd%?vrJwZT|5^W^aKO>mU$j^NRU0M!)oZq& z+DmP1_S8$`xs{c2owbkLUwh>|_MHZ8UZ-A7+m0Wxqt+k&RjWU{iiD|Gypy00TG4W? z=BgLaYnQTS*?G>@3!k@KhBnmij1CE_)-~Jc`ZxQPQ_a21I}N`rJEC} z&i=}|mbLGMXyG+`O-siz&42w?9%8Y0RQNpU@6*_MwK91wdxO)^e!3slGC!u|vi3Y5Y!Gp&C+VN$&a58NQ)z=v zCYSg)E&Zc9q{WahQ=eHPG`FNfhgj-^4bA6YiVs7(B9&wKNvD_f0`p58Qd>oWC9SX7 zN0E2htIs_p`{<5rfPSS&9sN1>qvX2Q*2=3&=kEd`rNuVrT>D5r7_Dp?WVGrKANHF0 zX(F;tN~*WT2fetgn{AX1UDk>{*Ope#g(4K2v6Z|2X3OB^myR**yLx8!Yd;drE}`r} z!}8xro|bsnpCsrMWCV6!MZeD$(ubs#b(q7J5)IMGbf4AGI@*8l*|*L{#N^%YBuwwx zR@muj0o|6*F<&w$EzPG^NUU|*G7fv2)m!H``>1Vii)fiLWP0zm_*~{^iHItqQkL;e zW2|;Jm3RkCGmJOB~)>6R+6=J+B@mQFFu-9qpNhWM}KJQHy zQi3_egV&O1tT~V=_RDX;4L+AZ z@V-)w*@Ag8;YkN8|03m(*}0p_QvKReZiJ%Sey=p*5&pr(_M~Erq*cp??NFi9C|V}# ze|BD5%QUHqIA-T2i&w-yxtqn2q$ORtlhG`ODV%vR88ciNCn9IozSKi1S!I@1k=)y> z%VW!j$VUy!g-n5bq-?$Z-YQzeIW%}+{3g!qPQ(0_2a&SuJD@>-C&wk1`_#~<5J+A0 zl6o2vRSQAFVpB5%>N_6GTu$HnGI=lxB+BPkDmsigncBIgf#M#s(y}O^q(9!LV z2}V@3r|i6>e@S7mR(zgjoAOp$(wKRcc*tr&mNZWGMXNr}LXwzd;q>ZB)@i=PW1oqB z3l}Z{pZ!Tih0WsW#vjVWYk$2+L-}4q|6F!LTD*8=o+TGzr;JFh@&ZMYUhamX#=Fb0 z#k;pH)2;vRpSjAb2)UHwm*T|8K8{(Tl*-HgG#`hykXx5+mZ9f35Yr@N zBMto`Qc_-};CCW8hi<2lHW&($kWz$I{{&pJpHy_)Dq$KqXCyOWAPOw2CQ?u%O^Ly1 zGXz;eA^Bc17BW?`mk&ntkcH3K0oa#*$P`I3B|pR4 z3Pxo9yf$7|x?YMPK~2&sIYs^ID0ed0Uw^@m`V&fm)rJPS-I97RLRhJXozR++IC;7z zp^!b)UyeEPnX*TV&xP$Jsd#o!l)OU4zAhm>0pcNBnjfr?Nhn#KPCo>Q3a%_qKWF4# z2O^kQpTH_B)+n=U6})+oZkL5$lILuYHRLq}IWpSIv%HK%LY$ityN#CS8yNu1yo_T6 zEltG2XqCl-^r}5wN0#~uX9&QL34#!&CB~R?i{fylmdU9tlC&Z^{2r%Tb zW#h7N<#0e!+tOKNti}a1ZWJ_wd;ksM-$Wo<{xRHQFJV&2*VBi}8)|te=M*XgI7ust z%sW=pZaft!OHn91BT7}USY*^iCk9kN=pmF5bS2{`jDu(y5eY&I*-S zLn4Mbfie_>OS#bzj)RoaX#0x^LeQODK(4|2E)5XU93zZ1$tzK=E!X8hbPzf_eT<72 zT)q1(|!HzUU!~HEq3< z2vQi4BsLf|sry4J;|QE2D&aB2O2lDcw8>DlFR7G8l0B9QUg07LZFLlK;ML+Fa4RDv z>?yC!vqX3|;`aiWk%KOuQdy1abQVY<-YbXzMu-!74H@A1eHt>)3N4BZqAdamkyR&v zV_77{_bRB5b1)+W8w5DSDDZ+bqt`nP&E_EN1R)=LPye0xAWiRI5-FrWqN8uI$ZW`x zkO`6thmlA~$=(QRhJ`3s*eug8AY&JFBcc2Kcd>oAyth4BAul35?_X(4tztU`Oo{BE zf{NAf^PweTy*xHivDqN>qNC91_G##z`@GXIE-Pgk6hUlDXW$C`lR>8T53H0N5I(h= z5nC#di1-C>%!sur%G0l-k|N9CI$85-avM)^2)2_p1W7FmK z3nV4QS2zPZEP?{~a?$eB?aJoLTnSV3=G4!+r7Ww#oL;*IdZ{&@>dx;I64mQVAhIkD z4p8oAM@q6d)+b>TpNl{#B6*mv!)qdN+y%mLt~-rgWg(i6rf=h`MA~$g`jxs(mni1tv=v{l^-3j6dB&1Q;}q_pUNw}|rk7F%1^SKbXR0!5 zwx^L6cv*`C=*V&Kl@yO46@Ei6E#EC=AxCRtC1sSyAfLh+H_6tWf^Ff$b60z(gW~je ziW=vaMgpZ3qR0{y2QN}UV7&ZRSgJ1&M_ecOx`-PNV}<%-(2U6St*N|G}l)sjdeNGhnmuoI@(b|)bCy{iq=3<;>=BNRGi zVDT3TdXFZ^CK942-c8|dN~-m1BMbslYXVj zqsSIqnB;ZmAA~jX8UpX~&N;#wDLDxs3Ii5TK_*$c;!K?jYqC(oRRGXk+xM`h8bVS{ z%)i{Iq;^vnYfyzp}LHMo|09_B-)=6ktpJl+?jq?$*p7WB#a;7_U`Dxq^bxNe6sQJXrH#ucV5(N z-j}Pu7p~KKViP%JcD^J+o@f$JmI&Q+Nv;EXXdK2tO99@CXtn&DtXLL zi6BWDX_w){STnYsexhzl(c`<4(#ij5&ESH3gkD)D;9cF74u}CZ3I2S5q!eU==f#Dc z(w&Clv;&b9_1$G=SX2w)9lq_W(yGlw*;F9kqq^rPQi|ob7Iy^fg+tV&FgLUSCX5NKBSK%PjBMLh%b}`<{E$U6vdDQAyG+@MA?H;_bLLr6*KRcA5P%UO?mGiAg9}U1W#8YSorIrOk5GMZggmPr zbpAOKB10-!KsPD%uyul5sI7*UPQ_^*wCN8s7OiXiMET? zViV~iVaF;i@o_|Q*fL2i(k?O$IMszWlOUJ!AYw#j{>nL#r_=#1e2jT^%6*Dfd9Hj;WB1sf7~Suvj>jUrG} ziR-F!Q9qTh*_qVo=7X z^Xn_Pl8YqOy0l%GV_HlUMdTVrOGhYf?K*k{7K^qdS`u8y6{NbG>k^GQ>3aFd~U*m!n-}7+!CN9o#fDU@GR7>C{h@0l3|!IMqW}MCnrv1G{72KNF*t^L3-~r z5cM=_seMKfYlKYjL^@C=sXd9v^qdv?-*{0%^Fm}NRJ**XustPF;W4oPMpw`=`8Eme zcID|#ge*0OC_@_9ej2+dctKq;cu)P5?#M;yeQ1x7NtZ!J>I4`Uv8_ICKqY6!8y|Ca z3URwj^Z>c}&`;|AyE^BxyK_)$I|;IJP%rG>z7wHUDU5U~);L*+SoE;Tmk>}a>?OW) z*p=~2CPtu8QMkmCAgtYqpzgoJd+zNfN)Z8DC?CH7)oVOypuwV0@5U|YdDAYmT$*U? zDGN@JNf)I}=hu4z&ljCC}NN{OS zW?K-BkL9)+=ilY2dFCA`gaWl+axM85k>Ejh{fywC-NCgVP&z3WcF@YZ_W1RO>&#fz5TKR_psYT}}F zy98rN#TD4Sz@v;<8|kc2#BJol2_vYtvkh<^d;5;w(zx5`a7s$I?fIBS@am5gH+t!{-X8U@s+ zOvwJxeL+>u+jbhVCY@VCgbO4yDmzO#H=j4UU-dqPP*BQo19aV)S-o@%OGR3gCm)8p zsPiLHuN5b$SGi(g25aUhIIN91uyL3K@xFQ3C^ zSI6IU13pP;JrQ64utN28P??%~x~!FU6wtHanFRTDg(+3=PEkMYq<8a_Ll9k5qJk*d zB>vP7j6Bx=NrBE8Oj;5^V3DXzblhMcdWheSc1J~b)(8bxWamg$cY3M%oq*VM1P^+M zVT*@MDRa~>`>69)=pd7b%7_QN)!sV$Fi(v}O7bnj68^l?3e$)RDH5d7?m9_~sqKh` z>5>ad^aE?@XH;D`6 z0o0eL;u!axg5urSbH{OrY9ESg)XxMI8_K}wFgp#YJPV>sI>0Cyn?NVX3e_%uE`%?` zEjT%myJrHb1Esy4_Euaa)eUCQ)#flnV z{jKCAfoGfn#Jvk3M&pUXrq1vZNNrD!BD*PRo>v4hJmuB|z?q$y5+M-h!* zdgLS*?GT$j27y*=B#nugPAgCW6eyY?X)vcxl9@)d_i>B2uYk-KZslz7Ysgr1efj z^M1`82;?M@3z~)KlcPOvKf6v?G*g*dZ4Vmfg@kDq)M@KQE1z>Ah2AX1J%iJl_(>WA z>1ehn^*a${qJ^CF-~GcS3H`xcz0g#_Pm-XAZfwOBRr+4Z@MWj~&g+PV5QSxgDwbH& zsA9exBN5i293%>*@V6Lm7*x~lPpX3)IV`|VaZPGsBgkxg(y-}2P;5XhoYb?P=>u4i zd(4(oIZ?&Evh*CsMnC0&B_v{zj;meWG=~NxMXmzCBGv-I__N!8CRZk|?PqLHBTdC^ z7hX}0nxtSEHlA!(!48?aVfE)_3B@;x-R)z=-`%qoCD;Ah-0@=L&aEK}WuZcJfBizg zp_BoRyYMdOz6dPIzfJa&bKo-9q;{akB^W{fzidbOSn=$Sxxgtv#bEyytQaJVR(|X+ z?nwK{ZpbGt_OArN6s~0%_&kxMCrv+14ju4WEGe`Ro)ypj^i~7cprY_Ed>nE8r0uWm z<-#b7QX}$qoHn@)Ix1j4<*VUKecRHo*{@Sd_noizl4t+w?f6f#IE4%p-?u1_loQfK zFX*6kfMlsYpdXPNZ>`x7bJvGnVtO)p-bKPHjh2i0Sn+J{mvE!D|@EZy|*j#I`N&|`B%8%P^)dlll!vl_VRXvkJE+QTTz4p zcj52t-%r@$PGz8fEO`ic3maVU{&krSiNDxY);9}oJGy~$k|>tBwejaFs18IPA=@0t zBZ)~uv)rpaOa4+&EY&Mi*-Z7pYH=#)`>XGjyH6mo zXzi4X5Q3)%@?*_ILAp^`96zu02ELzJF0#HiG82v0{fVk*mpDm|MV=C^- z&;)&@j-zw=tI))WNrC=*N;50h8qF_C;9S zkX;`x!GBXRpI;t`Tp1TXI{m)rOl!m`1CV9@6-B&w^)vzwR*VFfW1?_2p{! zdTuTt#hr`p`V+gOh+D>h>#e-LG`}vI^tTPmaP0w!kUE?t&-Sr#0|4PtO$3)^7Il!Z zz5$49K7$JpftTp$B$S-vn;vlyG)7^N4At*~XM6q9unXdPyndhlXfDUUNLW8Eq$8wj zR;nfRk3norm1{^ca@!VEh*j1pBl1QT)CTFJ8RNF(*`JcDp$9jK-zbdh7qS)YY5h=) zNv9m4V2Ssm4}i|^(#tm6~H@79twQm)f6mC0C1hBqTlD*^fwcdYpDXc*(5Zs?U6(E zm6{(M!07Z8)n;fzcaM+>113IJJQSGyGGN=+WAtaVk-;1<7?-}JfSVUB>s4wu!7|2` ze(TNs((Tq$?sS2d{Jdo)3!MVg4bGB>fPTrJFT$Jq=1CD`Zu8N4|0Q*v&e7b!n5V&% z@ON$W9QGMfs6P*e6bPt%EP3|#t8Z$7KMn|{g?oB3LV9X}Vl>;7)IH%qAqb%`ZYI#w z;6?-;#UYRtYPbo|KUOb-mplac zeaEu)?RPLn_K@EE{Rd$vEr|srydrigvT>a>$#+<0j6bD!3eKo`o)ypjK7VMRj=-2; zVPW(Yxj(W1ebXSjA`5`QmYJcLPUiwDpWQ=P=_2-<+4ME~DnC{{`!{yO=8-lpR{j!y zC4HGs9E8vl>!g)0i-iS_PRSZ4M{j{b@fJLlrf?I6Vov^8@a*5;y07%HA7CN%Lp z0V(UI88lHIGrbCNrDHjv=hkQ${7rY(2zHXP&xOKc{aEqr&ut1S7S!3bqVsb4ZzX8R z&ngnv?Hgm9ACk8DkU#~BT z28Yp@Si^3~1#iX~q`AW;(igC4c3(MAEw@`3yFion?~>R4%Q%=}EF5JCk@jmzSQZ!7 zp9^Na+b~?X$3j^k%1O z5v5rhz2+%mV)B*7 zVf+YcIGepBaafYa6fXQ+^AO;-F6P4up03L!E9CeJgvf;slFt;O{KhLfVQI+9k=2^Iih35#+k@bURs z@(|$d;MOz{G%VI1U-V6kKk!0?7K@r)m!0*~Va7hx#qGx}Ka?7!@V z!WqtM>+Q>ST#G1QZm<*3J1x8^R49 z;-Pbgnl8UP3E(B)=9%$J*X}m<%=$F2^IE6%P%;c6z8% zj!O&9jq*ffNZP4AX~;2!jr6dg9D~Z}`y^{(8aBgZHdg|(<%-UtLDpU@K1&`Vy7yzK zPNKciPi6T>=jumpPk|tV>i#Wv*p3LwmE?!c_ptHa?J$gtRsDuU}dX#^cJVcCO zL5Xd>A~ClQK71WpB0GC`oz#7t^PiDqpu9#Qb#-1;zW(Oa+fJ*OimmAYuxEy~K8YEO;om z*dAvwuP66Ex?Z^6(`Y`~AdTRzw*rn4B+!1!cp1v0?rCK@AQyP%C1JkmqmL<%ug8F! z^`{cMuh=#oE|y4s`gnwjW>R9{5QB3<8YYu97iHm)aDB`qGP%V#Q-<0RBg#j7EO}9| z|Kt;>q)PcQTW}-((nDwV<1VmOIQfn>OJ8i`x@_^)LMt^uYEpU6{+z)h%t_s)d@Olr zsNOF7zZRc3P%MQ&-g97I7?Ff&rp+e=i6j#CC;$i9GG9@qC)a+*$W{P`?ufJIAt7u` zXaiPq{)Al-hQOpet}hTV>HT)aVcRZq#gMT=+m4AV9E`AV*8*;T&~Y-!+pLfE&i0PZ z!asOh{@L=iWLjR!*XHCV5))!-c=)tr|KQ?<@#k0NLQd%L3 z4nPCBaQBs<_c7)P&nquDhoTcj&LI&@UE=~vy&@<=WTw|NXJ9uCSftrjEuU-t(vVYl zl@M5;=?qJ0iBmAJzk=3IwWi8rhtU%;-2=T!B%kCarmNN5aN&HH!DLt1SU=Z1G=%!J zAcCW$Boq(sx9Sg7;EYHbNZ~Uza-ScmqmktX47R#2df`f|eH51l6w70%^2d^ghS<+O zySzPQ3iGJIfSFQ{X~P^H(O^RAJUv*SFdZN!x0*9d%Jk)(Y$DmzH&{}1m6h=zqK91! zlaTtGA0|Sa$nWm7n^BOwX`{l@;m&rZWrhy8Z_aEYDB9XdD|{|_i12U4n-fRTtn~Jl zgFv)Ag_3)In|mNT&Vo+1!nWl~Nx7M-jTB0LbDC$xL%;QZpy36IbX`Yw&YtOS_ z4j;p$XEA1B-AudA?g5k*Gr$PvXQYREnIw`H{laI(Lqsru+Z3rpGt=cK4P@}Jk1j9f zI1)7$N*6JwO@YejzEL>#m!EuMoU`O1;&P4q5aH&M9CsGb>h&!TzRp^Fi{_gM260Zt zYD}VfpKj~)sTOW^eJuG)#Krg(iqYQi)%w9julZHsdT`$|?h(tyS2Lcl)h@5%+v%FE zA4?t@F4z%x;i4g&HWSkXd>=iSj(?akr8n)Da_Wtl#XBOvXrS7xU+FVGrX=hs8eFjc zUGb0*+vT8Pg+K&rrf*W<2HWE(h?A!vOP_IQ+!-ib1`Q*}AU{S^%Y=l{Bb3cU)?~Gh zB@Yc(1QYghhP94-_w?j3^4AY8q|U`}N&fb0%la$OWWnvM1k=#5=Q2uurZJg`or}rC zXURi?EkYP_Cz`ay`_GJON)`S1lI&WeYC@sU>;5r&qF0OVFUvD^uG<`C#@ zW3X;_BLX?ANwGr&l>J4)S@KXY!>@$+M+|D-PRh=xIyX~pED$3V7fQH`rGUPg?OP-( zq+oD|hR%->iF|qQ<$EK2e=K<@Sm9KD0i;p9_2VVOk_Zy-$~d?XMRu!&9NzKc6U8qo zuQcmaX3cyoc_`RN4?Q!(r;*>0Hc7=1(*Zw2+E4Yf{BRFJh<1jR>ZBXU0r#lp+@kQa z@_;y{2nHpbcmgiBeLbHQ4+Wu%HH=b^ zpmtvJIB0s2aNn80MVwX!t^XEWuI5vMBi(4*j0a^>3SVUMl*hB=At92lVG8ZiawLCE z8qMcl@r(o-z`->U$5{Sikv;(E(TE{=+n!VDY-r=QI1{+jea%C{ZzD{&=l&Dsg?rPo zHZ8`EeN+H?^VBZEfm*(uGOa0Mxb|#0^5KQ)$wjVta7FNU$wR=5bTmN8uV`u11Uuvlo^J)6nt%yW;QeB?7ca%WX??U~*N+6pmx^LB0Bbm|k2})X!o1^6VhjSM(2pe#1y=x7&W|ge@F-t@p@tOt z0oLdk8kAn^_7#4DDhde5SWzUL-hDc~F-w27cZ+Qb)*jDIv8qZLUI#^=>XW&c1 z{Oa0A?AbhLkXGF=TGBsBhR`-cdw-WaG*oZyFqW$U!Ilw!i|7D2#sjrxs*;Q+awAfJ z-R$B=ZdVSq^GzIzF;cvb4AK~UV+j*)ewRE{G;g1S3WM<$aLL_STpX00y|aQIwHb3+ zBTH5mk;!I&b(nc-@&xnaI%G*O9YRFbXT?LqjCbxq!}8M%y5v}0dbt5tzZXMj_~wij zX~Hr+9Jr2)hVID(OdgLYQ9MBx_!_2HOqjpZf;iogLhT;l*p^ zi*NFNgCXAi${sfECr$d+zCJ%o9vZgyGkT7}&B=oR{NyUeq7-!_g6r9XE0mOgf@S2x zmyoQ5ocdj1GJqpUOPVbB2EdljC4Y%9erC4#j(y`$f-;I^OzfV!K8WB-)@0;qjZGwp zG(*q~8)D5@*Kp}_hstqDzjCch#`i3Fh=5;L=F7kpQp)wj7*9p=H~7e&&KGoQLOGs= zrG(06T5&pr**S)G`r}IeUHZU`t=|Tf-w--YXsons>7R>`=`aX>g&9mxM^KH||WBGy{v3e1lUk zE&DqUWbfY^PN+g6JcV#qrZMipMB|So4-FpO{S)sUOvny&J&2ErOW)~E10GEtPcC6t zghdlmLKRDl5=y&qYb0)%*wCSwTsCbrYKV4Y=D zksy$hLNq=WN&0Rxj-LKJTZ^%RV5aXCPi*~I@enX$x8OyM6X@|ojhoR!9?*E>0(h~~ zzf)fb*V2r#6fV3 z5sRY*4cyCb+-%A5-rW4hf2I6BC>(c^?40!u1^Hq+1k4bE5K%_wU~hVPO2P+Ytr1M1 zT*K2D;^evaT^UFc$I4Kl<}~!jl81up<#lQ~U*9qi1MX>Ce{plUw3;70U=9ek9rq5> zxueF|XCAfv@X>fuy{d$+2v0wl8;NPh&eP9txWG#?f)#+n4{HkmgAyT1(fyZ#0-e zwDVM1;h?sVO_H86F^JC@0p(ZgywH8p!1x|rhJ#g$C8JJ*cG3`-Sc1s);@E6 z{e}E6hv`9t@Ds)G#sxzD!DO!XHhMD=o4d4)&br_;6UImjGUwt1=ob?X}LAcyMnMx8by=ZrBo%D0q{9@@*r17!fA>p!) zm$&f4B~-YV!o0dUo)qwHv9!1MAenft5m=V+Gu*`Xb;rw2dtxjXao5L+hlE%pzV*yu zHzyEurB}85KAfg!X753!qsQ;_J)7h^R=*CC{-cV@Twa)OoL~;m?jK7Y8m?8Yy#ox4 zu(o`EmKkgzKDshIh3mm@K85bgk_p$KzzP;!}F6y@TG7)mqOKQA4~qykb}D| zA4e5AV%%@E8F7K&!3bWH*I5@MZvIJcAwX5}pxwxbZ&sA2V4TU?A4?u0LV0W`!((&x z?R_Lnsa8C6r$K=rw+xsEvDUYv#yCvL3vKaZejUzP^3V|BWeeo7E`a(5h0Leqn$3Eg z#$cQa=$!B-yMVix8oI|7B(`lJWtau5-h4%l2km?;d58#h%{P=y;~Eh1LyrJ$pNwaD z(6DG35eX}|F9=JjgR016>~_R>jqj!a28&N;#Y4ii?F`wBr!ubmi)UWZar;a=bgm=R zAc3r={57|dY9tw&C-?9r!Rtx@aETi?TRs*%G=wN$LB@(y&?=YBr%8<2GxZVkbJK$X zsxt-47^SNX87~&;76$tO=U0SzorVhq+-UPz@({6pB_S85g2azY>4s+{COVoPQh5`A z?18dd-xE zc|{zO1$xF9zx1)>AtF4@+u_u(Lvb9kR|REic|W|jCj@XiWRELS9q~;YA>MT7Gghb6 zP+vMg_$JOGr*c+2G_=@mUA{honDc`%;v?q=BL?5xSBu-*j2Y!_M$qGhoLB zKJts6SH55TaS{xc zh$oD@m5HAZ8Wi3#9-+*!FSZd*P8V>@C71N@Tp=EO?;i!6{;}X8A-_`%MqzPwC0`1N zi%E<@c^t%NLctKYQtI!emq=#u^+9}daJ_P@^phuA__)LHyX4>ef0R}Ssd0e-0FvAm ArT_o{ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..9a79fe1e2c97b4f6037b478f1962c785e69652b6 GIT binary patch literal 4799 zcmX9?c|26#8&1|p_6A8&O_bdTlXWoJ##(kt$`TpN5LqgXWirYd4cV8-l908CLGdNq zFc@XuitMuXJGbAT_kPYj_dW0PywCH#=S(yUOVAT9i=%&0W(GGkrJQ{IkHpw@b>A;w zf4wcfcI2k$;pdj3!U4@Z{gQ#eT)V1j{RCaiBjLtd(o??F=?-~@;tFcp6}BsT!FyxU z8ES_c>jh`j4_9`ukN*CdGX8t;XN*Jr;AemP%}2WvKmP7-tnU2{Iar$e8+>@!|Ki7^ z{h7HlYJb-#8I5~0>&9mq4-U`#2yRu9l1(xA39C7)-lDTp@*|hCb4Vogbl!=jU~tJZQezqB)50gysPP( zi!mIIh}X6eDqfQGAy30ZDb`}lF+7e4r>ev-DH~k`nL#1PC6e^q&$X)G)VYu>ODKi&b0^cN-O7sur?;7-U@ zWWlMT6kD1S7utNd)GL`L&m?d!+?3jP5G;76+0}w}M?n*=t=YjI{vC_<6(u^iVotns zHS0o*3OW zH8Hd?m-bxJ;m1tJk)JPMZp5hO=C}PZRCZu+<-Qy?944}aDwInqfp?pK8}>#KU6uWp z-g&Tbvtum(STJ$yBT_+{{ho;0xxjoLb9LAmr`L8x%aA_oI-j=nK;~xWuqXSqgBN_c z<|p;enkTCKG2@}m78qOwnb%J#e2U%}^VXOwEa(M!cT5(xw`kHZvoejyc-7X}1Er?s zw}VQoYSfl%Cx!<+x5VWvUd~J8n$sm$!0lsZQT*(eg%<0qW{z(ry<9I}Ctm!Ni=p{JF09lj@-Mn+bMw2)=-g zjwEqvJKI!-Pdl=q40~$EJOs&wu23oY@T{Lvuy4Cqd0W{C-&?#iFksR5VEzPl_R>J* z!_mcmi(>RAq?|l=0^%0&E2%z7LkUis8-QvLN4_Y%=fBW-0-I}>8ynOp!Qep7C(yzW zS8YDW$~7Ci3nReE$P|>Q@Eq!Le2_kt5GI==q82j?(4VB!CB~?o&u{;Q&seCAsfRft z71BC3om8&H?ZU^OiIakMdp>jX;+pV z@V!Jz_VRl+{yuF3k=GD&*{wT|vTNI8XZS+;pnCgwO;h@q>DfJk+g_>=6n86TNG#<$ zhdM^K#j(tHN}IoHt+MuO0&}{t5P}uwvEYoj*zCF}sBw8Rsoun8@Tlts8&umd>+}up zlcip!+v-W$&9xC^1fQQ8F^z4VzTjG$IR=hvFWCp|v2QEGrKGBQavYbeEF#8jx^sjhiu0 zqad;EX|B8{_$F&eC#N-+;PA+~=e}oJ6sF~9(Ot&grfk!*k^lt1^fZlZnwrcyJo=+?0uX+j=ae#u}+Mi3cBoEAd z`Kke7aD%vbS?9Bz4p@SULy1_WAx!Y%?S2VKjq)#U6shB8k$G}bk>UWQ^PXI<KZff&mxaeE@ls*WamzJG+c-qVK3ixh-VunM> zH^l4Wx6^vO0~wcf_`3{WC(22bgsW!tBZmq=2}DpCL=DB6dt)u+A4qP_x~Y-#ztZ|- zKhT(&3O{@!%=p>Slox-<^IqDc`%fg@8(6 zKo6O#ahgX^WJAy2$>lv%V&1Soo~yGwisw;LL+MzoN&`#I!~kK4<+@|&CJ)6AO`#wa(VwcpVHT`PjJM50c1P<%3f^2 zpZKU@CQ&=39)MR91Fx6tFoWB$SUL290iHC9T-3XRZRjd3mv{$Q-9@ zlTeDKeq`xE^d5bQ^=9)J7XC@Fqa!quvKpggNc+x~w3I-?8uuQ>RSMtB%e7i}(&Ec* zS}gKm{WDu2VoF+|ayXW2-6#D#0bAT?u;o``8Xj(TG^~Fh7g~#!&4WbpdW)x6KRZA} z8Y^ZXh0|}oK%cbfk?0-(T)ZX7xFcZAB2u5i6l=Ju2Wdk)dnvDBR-+%Qiujk|1+8Hy z{UvwvitP)<36oAe;uB>MXDQrZQ zgp&p*``F3wwzhyMKk#N=@)~1+!BC4b}02_d1%ZXj`AvnR%% z=N0-d^l;<8-2=GthK_3%;G59mWMQn-IDjX_R?>CcE7#@x*5Ie&s~}5)7QI@GtgKE; zcDEM3dqP&7->Cna|5DHA>C0n@B#v@7IYioSuvp#i{GnqZ0Q1ZRpYC$w7H74x_8scM zn@)lZX~SpmgMUzLvjC5vPLjH!8A$fK4=0zi=DU-hgdAc%@&NA(ilgzf5Wz#ora6`czNyH$78K*k}x$^zHsH zZ~0QM@xA41_d-_8^+O;|+toTc`=NXz@RbOc7*0j(ENLKa*t124P3vx1a@=Sn%8SG(A1i z)sz=!m5{#0ZYR(7@C_9!Z3UuojC`(CHTdaguwiKaol=20sQk@$?i=?EsrgGpe0={U zEA*4t8W4bQ%Ay+~;+St!xk_;gl9fie@3JaIE^!orGsyI3Q$FysLbhR)> zK9P$E5`am^W?I>o-*itOBUt^`>7VQ>PN`g+3SF<%;U1Ie-(A=4ijg6eGns5rak{2# z7y7~>1Du>YzIJ%4G(6tEU~g)B7`gvzJSLmYWN_qp3z7;L6?7AFhEu*t5?e@D++sw-(^A^M*kOo?b-7|;+kG-wNR z(2mVZfCTDr#LOV>qV!@xUIqQ@f95?N=L+M?x0V|~-6WTsazv0NqZW!jD7g!RX5ewe zp!hXb%v|}!i{Gs&puW3Kr-xt;^wTzzpr=)OKc~(f`}_jZAnC3M1mCG%)&)1Ik`+yU%?W z8)7nnF(aYof&^z86l;=vWig>Azwc=bc~wWx&*s&81al?b$*Qys?C>faqJ<3xk%MX-hEn?SDCl;c4v= z(XALKcrv*7< zMIVLa=G7mb+*~^-0yq)!a@1BFax0nG*Z`0~{Z8^UVR5K!{phm$qy4VJ_+CW!BcDxj_JEd|CvM9*x8BrV)JeIh7F5_DK83h-~e8x@EtGT6&Rge%1~>e~v1w zPH?T?Pqc@i=@rsnvG|UxU5ht!iwod4vyzrb347#nIqzgyoKcQ4n5?L!MN#q&=d5_~ z8!D%9=ZNU7^z!j1GwJ!K&&?bpwbbrxXcAptbM&cIHCaw-bJa^r)V*jC0X(*tzJp4x zy}Ces;M8){z2;RHfr>b`i#&Hur1#{Q`oP}qf;K<<+j-48!=>X?zxZ`O+j$THxBVJ? z7B2Qe%P2j&B^+4cNV2Fy^o9kgL$3}RLR#v-vDWJ7DbzTm>vTCC0D;bru35BFpXB;} z{tk7Jh{NFVDirm8fL$n=sweSK0!u3amim=gGc< zXZKd9m2f*Zkk(fQni5K)Rn53l0!uVIhxd;`vul7#ihB&E(-!Qc8bZ|-$mgpGfw%iX zJrbPzm3Ibs7ZssaP<#36%Jd}*ml%l`Hol%dKm)B$lw{ovPv=eixlRfPJaA7KWhO7a zVGN=7Z__F@1=ANJtWt1HF~gp1lGJZ zJ<7lb&&*O#G={vSp;G%Em;77XXxWjP{s2yN#mTU&tOhoFDZ~hG2I8M9D=wz5rDFs0 z26IPuGXISoe@+iC%CTiHWEXW5!m)Etf&L{@ETvoul7@O=Wg&y4R8MG<8K8*%ST1#U z;m?&4yj&+eijr=29a&7hBRs^s|3_fwd?1B9hdBl6ukGGKwupJ$qnUICjNvD><279| zJai=ToRIiW`sBe?`PwWmzH0E3T*j~<_!zEBByXGaPuv=8bTw_79*K0J1gfFN)>#>G zh;&w!782uOeWR?}1!Q5Sr%F7E2_4Qt1!94 zA?svfbm&4}Wlsc`AW@i#ZloF3Dhq9{-kN@4Gc-sMx+9pS<4w`Pe4n`q8~in4%_ literal 0 HcmV?d00001 diff --git a/tests/resources/quandl_samples/MSFT.csv.gz b/tests/resources/quandl_samples/MSFT.csv.gz new file mode 100644 index 0000000000000000000000000000000000000000..b55a4ec27943fc6bfaa2cf1aae98dee78f26be86 GIT binary patch literal 11579 zcmWk!WkZu~6y8R6cXxM5cc*}KgMhT4FnTmdgOo_Uh_rOqMt2C(-Q6Ym&L8l6InTMy zRX0s6Dyn~ohXLqSQ%^~s+t$k)kY+wIE-2!*GNCt$5b%u+56Gd3`&_nb3ur(+BvNIqZ{QRu_`rvN)dOq?R z@O(1S^|Y60`Fc4r@p?le9rV2N`f?tX_wraQ{d~VC9r%1@8T4{c)%98|{c=-m+4-`o z{rd9k|D3aWka&|oc-=p=Upf6?uDTGU@;6Gq^YH#xd{uu&MaIiJy{e)g*EyZ{mi@sE zflcpBelb;c5=;B2x3FXj#dl9M7?TB>WHnsY?YLAnChv}lq`*2;a_m^aCu=!87+qs6 zi^~hymBOAm<(h%Vb#wZ&W=Oy1?-}b=EGB9%iQ0mP<2GXciC$IkyBJ^vj zYte)uF-JO`_ucuo4neGH?g~|l{^Xgnv_NPI$8QM_+73_w}mDL)M7WbT z>C-J_FzNV?1;G?d*#su&0tK?2N8P+K?HNAr)1*t`R0at-Mt$y~TxLV|&l<#k{ju|n zeD<5=tikY6U+_)3QiPO{raKoGjP}rou-((|cWCl^4UqJl6I#1zfB3RWB!HQ!Y>!^X z@;XQcvaG(c+ZB|$6dJlzdt-tVkdvR&^2ay@86@(yprk1NgUAHy73@2cpDRPzWU zI45eRUZ#dwpNfEA&^O=NISoD6{mdg7^QHF1cuUCn>E8LMD6z4Prq$e0aB%15tBbha3}02*4)iu7aI$wk2Z% zd+Vb=a1|rTF&!F){8}6VLYrVwFm%s5Jh_d5GTU@BX(OwUy@-+j4%gU_gU9?#qNa*vlXb`LVu2_eeB0n4t|B+AFyXu0lUD3q)+BZ1EqKI<52{F4`XV6idLA+mG2zN z^;3Kh73_@GW7xz3tRhy00*=o~W%&_t+bWE!I0C8><*S_^(IIl?sJ6mike5+xXqewr zJ`ABWS#YESaQoa4!=7D7l}S^yWu7>~oa+4Z2hMo|dd_hZ=R8U3m80KbUHHxdN3h{L z^i6AR@0iccG6-KeKagU5WvgSM3FM4;e+ZtsLlOwv2}S!#(YMZnQvD+c8)7o&|7aql zo|^Y$eMIyLz5CaPjo*XS=N6t>)KEyApRt1K@4`mRRxy_E(D-dS8PW5#BqA`aAd37I z#w32S3yIa|Uc6=ssLFu4&I=M037dK}KOPMat#eU&>lGCR%Vi+m4waVg!>CM$)>I(x z>`3+Okq|<)kdZiV^?+KKqUo*}%YkKB0*;ZbfQZT>^!K>@X@v+89fpyVH1a50y;9{t=$!tf2 zct1RX33{+-(e*7ZNpt?FWPO0e+nv)CUqxMt5)K#B8z3HMAk3S;AW~AzvPJ3RS|gv% z5+s?^@r05*a$+32eXha3!hqr1bkPC~0*Po_%_waEFFGV@rY9X@OLwyT^Lx|R19-^A zm`USA>EBOwf@z4?Q14Nau-eVZXgG6EQ+7jwcfu_ge#67vY4ptysGF%7%5NcOV(pF|j)6#XOArEY_xWI9oBQFGhE|IcFSVu%x zyL}+vB*Ue(3FeODB7s@)3>2QB`h`i_BN5yvS#=DD$+V?k3r`XKnwnPGVhHvN&UXEn zPXhnVxgE&CdM{%Q2#FZMGIdwUMX&^j4NG_M?VBv|fD4x*pNaph4%4pIo&*G`KcFW0 zbi!dYo8^atv|4~jXv7#y#7J~d6iiB^gro8Hi-d!9J&@W-)b~)wG3xxPg{D95Da;pe z*-bH|4oT_NdhtmpM+X`MR)D5?Pc7?%eB>#u0S_ObblW}gz{oAl3MpDuuG02xj1-t< z50h9<6j3udK0e^XZYzcrYE7LtAnx!MU47t*p}&owREGSMuWK_+oy7-G^rC_{$q=s0 z)(xYbD$?Q&4?z7y2qDWkz5?H@KO0zQsDV{0^``Uk{A)7v;@gK zE{4Ml6aNs2Qh)dF?gBr7lwi_DwNis#Z!urKoe3kLHtU9F!peMT6o_QVFDRj_w!qX+nAfi$R&qgy(bsnN+14>{ikrl^}mW62}2-TCG*CfUS z=tenQ=J-q`B?@o@uvXVtp<6n@!%!FI!2-mop?2{KOwg>qh*41mB-=WMOUSDkJ^G24 z(vpD&pL*8>u|EAw?ZoBC{&Mn!cPSWKm}L((j#z`;Q)x^NF;L_ifSqE`5yKY*6#beNjgZ5N3SePH ze&u$PLi0Y{uoIFoF}2v5;C$pmxlG}eX>-WA&c0zo;PtWfU`xTs}zpQtp8jDnJ5WEv$iJL zqUiyy^wXf}eRJl;>B7Ujl`<-A_p&Tz?=D_Yf6fm`jreLX@U~ScC=oOIsNjr62ec{B z?iO^kr^tV=zBG0+Q3xxFh}*Iy0hHxj448nb=4^ph`o%!|Oj{BqRhsJydNL5W*SVMs zHQ0k1w^If|R%&KUMeFOKZahosng!i404<@ZPk`4SRGqLr(<6B`0haOYH?MiT(43H>+PAuQp(cl=W%{xizOwDSGhogW&iXlY zu#{tWewRfu!Svqf`75edv!OM#y1R`je<3XplsPA2_p?VH?Sd~o@07O-5j%wrGB2Ce zjQLxV4B;5}i?BB!6nKMPcr+vXjN2KhmM8)C#-X3|Nbj#Z68K9?;3zrqk%8=~=uzI* z*v~T1Y_x#=DoPWu{T+~r{BGPvXTz4G9KTFme(UkSWTBTfdA40hb&TqSH~hTqimvJ* zv@(Y&+TJ(RTm@t9havyytK>n9+|>ceVA|)^dTWJk_jReQ;RSLAtCJm1B7qlMGU)hT z^Ci*w`U-v-P1_>@y$3U30x*`(;Ki?&4|t3?x15S8d_Pl;TlN(o;CIv&DPun&=fR`l zA@3vGG)u%q#1;>~;rv+lz!~g)huSxKYth0L%~-Gs*J6st&FW0zgB8$qqFXZ~C}|tO zqhSp3%{Njj_CW*$@h3-FbnGJ;{iGN=GTyL$(n1uRTu9hf4Zkey7*j1Mw^}0%#3KI01Off zbw{w&I8`J;5lBtbyT?Loq9r|6Zq8{F;_no!Yj9pawN@bkkhoblXd|0=YDFw5Dbl?# zl0B^e8VuC=6bt`*wa+}`L4o!o{L~$xqQ<2obg>>zOi8GO-B(M1td5r>OU*wQq*{4c zqF9Or4)MN!DeL*uijtp00}`e@*UR%3SjT9t2JeZdS(6P<}JdemrXKtXz;>3j2#Me>&a z21HNRJ+C9(z&Zf+A&Yo(CoR$pmDQux{+~cSk&$sj{92TQZp??fJ>C@w@%FbuPVCF$ zy*kESFR1#*us!P3k>;eQ4y?EU5xSZ7A3VmbT9drd^gJGmXv12+>1 z6o_nwXWWczh0wp`QvJU9s;`52&z;lV=hkGVNy$deh}3{6`ooSQ-+e|Vz$fqz0){$q zSUgu3Uy!4vIzd)xRKRI09TrU2-?vEa!0M9<(S~uD-Fa0=&@2}NL>%x-Xg4Dp$ImJ^ z9>$>#LH}tF7>f6w6M>_#LuMW=0%(r6|M*G}_D%+x>=^2H(G=ttIlhTIs9~Z5*+hFh)JW>hWd>B^&Tx3dXBT3N7{N`Zb^)&o3(e)D|4zgl% zt+^_vJDGIpCi1%LZbl1&D`s~nN37SM@>>wER&P2BXkM0hQ(AXryxo)>We*^N+U36c z0TqN8Kue!^3(`u-00@~6*7-g(q`zW_J1J54s)tdtdZ~N8k}DV#58>wLFhTkPk3KKZN|=O51MN5`c&|sF zkRFYq?af%-DYlUjC999l_Z z`>l{o?@NC35Cpj5dU0dhiI(Q)*Mva$(h1btkT(0VkQ&J%c!?{o%#sJR%6q>#J{WF@T}bdZGeVx8Bc zGTYVx_eI_ST*q=d;)cE=E=0GGXt+p~hVD%1vA!83guZHm-GgvZs^5M*`=X??^c6w`s=R5w{Wp zUA!BqOO@<_VUlK}PRn^YW6qY#uTH#mJ^Aa9q$sRhFel}3v@MZ=p~TZky^^SU91s!uOS~KwkV& z;B$hwu{=10L?P2Tp&c)bL){6phNHMuAr`9EXYPU;7x#Q^dKe^-8TQQlN|juk$$ zCQe6n#O9fiyjevm(O5Jw#Qv8c+|L(9CWMs_SQE>$a#jSU0%UQnCJb0mWayNF=V4JIRT#?;4GC%s{x{=T=yz@&&G=EZ){0* z@eSK9y(O{TrAPMr2kMVjyuQKquA$~=lPNnC7#?HGt>IZVf$;f}OxWTJT2Ffs`~-cS z?Y9q4-RbHL16+Qb863r$JJcd;p?-3=pK97xpiN$X7!71s=YOl%-EjN}fFC-`t z3r!L==)Qdkv5A`x z;_h?q!jqW4LqHXyl%B>jKs=~E7(TMViWuNf&>ItRbBt5n5p6;NfaGV6mU%zF+1q(T zo!1OU4?l$Zd({l;BD$sR(tdIT?Qv4+(rfCP@Y9%Cm`!!lWN;KsCsp@~IHpr4FI4Y! z4S61{0OYR8u%Og~k@4f!xx?oUAa|p1P}EmdHE<@Zrhf5s+1g78vX(j_nzr`fTP}1w zOB&P>X35+uJ^Zc8lX+K&8f@@iY$ipRqRrja%HQ7_mo_NE{j-UOY?l($t^7PsTEYrE zv){`K*%8_IDb|4;QzS$1vEf$>6k+?Rp}&)@7@NJ+1Gpj9%}&=IorHw|%xoL7I@Y*x zo#jM(l*h#3L|CQBJqCG%91!_tvOCB9%?g)!_C>euF4nFxqNS-#PNw9k`hZJUZAo=> zTtJzuqbWK#>jpQR0Y6rh#lDmBVT4h?rvW_?7rZzHeYXx$uTf1{8$F%p|5lNEp`T z%@ClTZ*E(R&uak~H(fCUBxcdQBpR>S%C?dChZYBt)_@)u5s@COWgqCt@T1K=1qMYv z`;z=*J}LzJ6rz2B1@(H5;`x`2gc7aClzkT_GE%9-^m>27 z`c38tI1&;NsZ8m}AnPY}`(_>!#|YBCqs51qwDSQ!5;^Dzp)DFu33l)y9MVQ<+DAw; zh*eQw>?_6{g<*9*<{>wp7kC}UrTfPXPFRtvMJPQ@-DPLLg%mZcB9(>Y*w86Bk%?fip!9I!#}R$x~86d$^j89jA~(w5w=0Kn}uoBT@qnOYsH=4Bq}lCXaYuFr{S zmN7KVE_@TNVB};gz6^&iJgq1k;Gn`5ZIfbkFB=s2of`aJs!A%=mVv2sfkkiVM)EOI z{0^GXP43|!aH^j0PquPJc4C`;ipVD>_5A)W7yp>S&4pl>8!A?|G)U-|dv7{X%^_~> z7cpspEebyg#p7Pc*SVP`=U|=P7QKr5p9qO;??q7lJ2%qYb8kkR6qr@9j>j*tRYc#d zG`#W4f$Kvoa9Sp4t&ln#KE#+tX#D4Qt)9 zIfB=J!37F_ppcKSA3_75*D#w|WpAZoNd(crSJ!jV%3EeRp=op}bK#{GZL{ZEteF9g zGbnvnd%}%+4Q0scY+BM=O}Br^e}RJ>&8VX+Qi1>vl-CSR z+4sZ|R*9h?zexH@^Jgcv^9#;T*%Q)zA$+vBbJ}ath;WLAv{_=%?5nBlqYEGp&(G4l z+Ubj`U!~#%Alf6Lsr2K&)B3B#e{V^5Zvx;#OJ1^@VsoY5#{)}OLVp5v2GXOut^Gvc zYD9Epc@--auw>iXywOy;KGL`1u#TX+5u^gNLPTtN3upHbn_z3t4i;~e!+FYe`NFSI zu6+ij!QEaQ*~^Q-B)CU>phxKqwZp^ZSmUmQ1-kJ?E1zC8as2t2rM4Kn9`dyghE+c5 z^>+8~U`e%aK8G9CKi_b2-J$3JA`CPmot*^t`xC5hd~C4_)Nd?9xw(S8U?eXZd@5;)R+X+MEwhL1BMw|_=3NSqWKVFBispasiUmi`-?V%Sf9 zp`9>h*X=2(Q`TxM=bd827?hLs-*^TKZm1ijc97t3(XTSk&|`#RZaq2byxhggu4GI! zbYB*}-)XE=?3|I`FlMU0 z3w5p=10@$9SMze5=$R*n8rgY6QkFq&hro9ii~Z23m zGl&!zbR3OGYs}tHfOZ_@(AragJ@}x?I{VVAO0rz~S7oMpA1NOgp4an9} zm+JuP!8DV4fNwJI&3(X|fUZt3^E71}zU{iCwQdlXiVo(>rf2|Je-`RgEJ78o`3VPf zBze|q>JP(diSo$%3#CT~-Eg#_<&IREb~sRb3e)?~(ANcYYs}8cc&k`~gke#!)-s-{Aq00^s6qG>pn{WnjK%)H^HXY+*-UWm!l_2sKiOF8)n=bx#;EbKEIwq4nGa&THQ2A-D@FrIIY|ga zF1L5wqYSN|{nF2^YC)Majmta10!(PXNMgASYqPgXV20w?Qj@NnZ;B<_kD;$mZ6gr) z&?fot>#V@1sTt@)^9E-Y8|XVouu&voPMi3KZeZd`cUXRA1zJPR?Ca2M3x&(~d00Y- zhU5cg-FZl5puq!#K)VDU_Zhkj?#W{kBj-QDXkFQauKl56mdNbvt@8?uWE z;PIsM@cBFdT=mCvvp3=EwYI275$DHa1fd}DF zt^|NMUvb@y9YK>$#Hv(Zrj%?R2I&+k`kTv#f5NLfK;5x@G3J#Vhm9yV@C2U8ElxO- z6F<@r^$EE5g!&8k1EmAT$`9u?qZC`-w%?O%Oq9vSc8v6dzLC=kav4gY#2M}Z_6$?X zPwrG3cZnx#kXgue^m@NtXsBUJpI_JFpLZTFo)>4n+8`0GBWBX4z_p|M62(p8n%`|s zA@UPHwFak$#>FChp$XQXA8T_?f@~<=+HFO^j7RIUEp>yGK7-^c@!eC^>h-wU`{#s+Yt;OMwL=Pw$8xl%^;MFIxtpR|`M}s;Cm0 zywty70~V;YRRL^s8QpNpY85e?B@WP5Tc8?Jn*Pff(f-M04hYbrCUB|`Gu$O6!~v!| z6&SMbr}Dg1AX(#y&a9~*#hX%cj)Mtz4oV|L_eNDLpv>XSL{$Xr4+WQ7taJ<$l6@7&^!Rk}Mc z=Wa4c!4|r9V-x!=$_n+eDZzUaM;XLCuHo;o*tpN+H=b1i#@sl1T5eQ#8(I+q`XqHo z{mrp<=Uwv$vz(ng0Y&DQU(dk5ps}>fklb8lomr39DI;{(d%Z-i7(>FE*yd$ciznlR zKTN9I?g3oil825WK1uKPV00TZHzK6!-groU;;gS_g@JS;LsNQB5f z(v&XwH5}Rq7rvnRk0Kx>;A22Tx%5ytD&lNZ_B&0^|t;R#_<#<}_ zK~SK51DDjiPpi0iT!f+QH#LP?JMaDC3W}cPHm}zz4U4NKcy`>?KN2zB(DC!A)$Gxz zbx^?lFp-d!L_uZ+cl3{WR^>3NX5Ibak(ji_h@STe#ur#H;>52xdSmdQGjaMSk(0Fp zhumqo0JgqzFqQAH<1wM?FWj)XGW}c69SV5=GQ7&zPw4DLbl&!d*ZQawh{kZ(4*pNv zkv&XVPF%ZNG-yD?!1z6x7_oRwMGA(SrJba`by2}I%$v;KyE_YKSA}4KaCG^#ke*&vbK6^?auF^Av zy>A+Iov&GD7RVN8*L@;$!^~J7j^nt-q;1U%C8uq=sDY(B2*M;Zl%SPx07{t z0M5KDj^|yDwBSPbnJ=Dv?KH9Mk%8Eb^zG98SBSEHrfr?{99Z3y@5anTD&@_-!Ze!e zm`oULxlIJ2q4pP;Wx7b^Nyf)>+-`yiB3m#KPEYtzJ zt%;5!PW&e`;idpo5#3(7hA4w8*Y5T=CgP5@t@4}&nC`5IvZl%j6^2iUYozY^&{#X)@4N@yHx#NQbj38^b5JaCn5PAWy;rVu zy6c#0QA}*Wr}))cE`4{i#M$?NiNLTBlv!g|*Ves%g@+^C3fi#OCvk(&CL4wnmzwAS zu9msA@Wr(nmZRUfLz3rDCaJV!PeS`eLM5jT^_cB~DClds*}Cmwq#>wmBI0&}G!i7r z;ieC7R~<6M+GME~F253cZ5Le?!o4p~+6Zjzs3QfE#zm~3hhJYxP}1JQ8Jq&yV|M(} z@u)N{9}0J)g9Q`RR`d1Dnwl%^E>?M<+`q@I&kjMLp7yyCT9mLfFpT5s@b?hc#=m({ zzL=gm=0s%uIz=?%L~dhJkf7}6pr7?Pk+t~m#A%iJhI^{I`E9%pawE*<4DND_LZTQ) z&QW{xmX76yZ{{v06Z(DmFlz~NP|quc1uVQV1|;Ei%F6s(@m$V37I1e} zZO?A%xU-NFbYE^#3bi?>8L(>LTWn9_L@M!H2S&SZb>XOBQKc!qwzRHwbU(^Myirw! z1gqyWb!d0<;?&D}m{mwO=OKQ*rtkAX~>P5fdFXm2uRB zVZjN{8#`q6b1|Q@eKv$}4(3Fn@%9TaD>HWDe{$*oUw^&l>MGNfEWpPzlPd^%w3bm@f%kjN?@fQBcdI8P5!v5O&P63RT*cdRdkOxRZ zq!nq{YIl~4ALtctN@^hROTk3rOeuiJMay(%>obMQji1)DfHgr~nPgVP30%8VAE*bV zI?9{-@6o{DB=c`%}2^XEkSCE)H;$X7ROX*{z|&)ad|{Ded4N*26$FmXnSpp zA{kU9XIpz@{BGEctJ?(dZ7E3nAMI}^<^4*=PqxwVeZpV-&5nytrBW3SLZ%+hC)xi_!RVQBiAH6rZGDK~ zWkH4YiSIM9@B_L%-8*BP;r$h*8NulQ4TPEqrm$jJPOCj{{N#pl$&9l>R)OXzOHDjf zl_CqbminQPV?8VVJCRLEJOUmS)RCrP!TCP1B{v|5iw z>0=>SRfh1*pRd(3#r*I7%X?9XLLaq17|J90Ph?d_C-9KG(&HR7ciL)7=H6o66(G6L zEnqA>oioP*C##+uo8c}9?=C*($|r=3PE{xYPlZzbAwz{r?8BX{v~R?Bh6@!G3g!9{?3-~ z;9zl<4S1M_Su3E~&F4q*Wp7SzC!%7BGVWJ`k43bCF94ME-xUm-{bESYcR}K;svzVF zjJecM8BYOQmmV4=o7;@dCwU`h@eVDJ#NTmwD6feCiu3uV{{q`hIEGiU(ga}?Igfi_ z0kS=;N@KSu7cjY5%sXXe++3%f)7Q9(>?uBJwC8Q@3H_&(XzCQqgI0yTCuD#9N^Ry_h4z(Mgozja z@B_M>DTt9`R1R&U-QIoNMV;A$ISNn-*WR!Tyq0u!CHZnEWSAx#Tw00ZvI<_o(cDFK z%o@QFuWTR6VuZapPaJ=W$Od?1Xm(vKMj6oJ1nCT;NL&cy@pj9r2XrgJx7m;P24L{R z$nQT8O~Ul99wZ0mrDU{AiD+w1=I@Xq>G<%xT^Q>9BE+NC3V(p%*}K|<2(IZkun)E% z8Q)zlMiN|kKhkJewSPv;+d-Gg_042_)U@$BQ-u2fUi)N&Tj_3NU-f&uH{L6nvT9Ut zac@}^-xK%>N)A^AZcb2W%PLP;!26CD5z^{<^mmgQBKx#Xl>-a!%TV&Bb9%{Q&`0UWM}yf?T8%1{X`W+klCNwnvId^q5=i&mk2rEAFH})peH8a zG-_Q+B%h0&{aJtC7R$}Y8VN|qH_JMnrxKLZiyr)149yeZo&VH7aQH4!w6%#UX7M+f zB(V|S>apR9(-x9rX2#xycQPD?rWG=GteXT!PfY`2K}I3NbqDQM_5PvV>H@~h8ygZ6 znY`x52eUU!@IA_>QJiO>a`nqEr=c#`G+G<3mAnApyDMLDVM@XNfLS6#Q>dMXV_jXn z&tR4_NK%efPQSF_Z+W7zYj%uZ#L9Y>jcgs(3^I0IlX6(Xet&GVPuMm5>w9!yXDEo1 zPHoPp5ZnG<-H_evAG@b#4a`;?&~fc4x$bNz(4KK|A$_FAaGbuH(_M%f>mIFf5N>cD zda1wo=Ymm9f>Y0iq8?fluZQ>WB#R+8-^A*#oH`39Vrz^}w&MBPexC^A6Y|`4s!MER zP_@Pd0PubiPi8hKhD81Lip^$oGN5@DM|Y87g!_gY55MbPf>B=fbVuGAde2w!7G zzB%w*2uj-1vDvo3JM5B#SF@AhDaP#_0kOxqyljSN(?q3By|~snru%kI3#ZBACTyYA zMvdc-`fibYF%g)@MCd7WGC}&bh`9-K^c_Q{g%G}D5h=~;zq=M&^wmZuIT(JQV$>j= U@lW@ZjKO!o7pe_@ssRAN{}^-OX#fBK literal 0 HcmV?d00001 diff --git a/tests/resources/quandl_samples/ZEN.csv.gz b/tests/resources/quandl_samples/ZEN.csv.gz new file mode 100644 index 0000000000000000000000000000000000000000..1f5f493bf89cf7fc39b188285b72521f1589355e GIT binary patch literal 3007 zcmV;w3qbTAiwFpNfpAs=|5`;(E@N|c0EJr1j^(%w-OpFd2e>XI>b(py0kQ}XAOW%; zpbJS2I@1HB^Z5Lbhq76g@1++lxU6%JB#LkMPd|Tu-~RmD`>)#{KmPuY?N8tSyZ!#_ zw|~8FfBW|J-~W8y{_wxIpFaNc@#X#3FWX;#`}*JIzbOB1fJTZb9?GK?pt3v@ zKlj?&ne+x|q+^R!e$XBRsDT#@Ei~xN>EjMc0#sJ+kJtoTo4gri>uufV4bVuVEKM?` zB%maZy!+H^KU2CeBUI2FiY&@-xwu3zu8bbeD6wL$`DQM9)`^hF*UHM&UB!yS|f&|aW! za-y!X?;FKcD}CIgtZ#w#+v5>JSCtD}%5~YLk26Xv7V5N&{Ag{yppp7Y>EVnLwiR%e z5(S&0Jo8FDTz2W9Y401}(;qQg0ku-cCcG*ZG$ z-bX$Q#?E3&v7Sk9fJTbtf?C_|fltOKo!!R5XogO4$i?HlmxLyNC;9N zkVC0kqCD2?@o+|&rZ*MsfQ+yOiPNcEwcWuPrK%T}540an--LSMbb&^y-a384&a7gL zOC)z-Mp&o&pSn>XqT09CF45;T%II{W9rB5FA7i@a8ZPi4)qpms;TrM`rzAdW&0iec zq-<}hVfm2Eo~IPL?C~YY$Y?h(9V^VCN+KFZdTNO;fkvA48FnZ@sSDwN)qA|;Lk`X; z+nktVK?eG0;sMLW1$j85>_{Bor!oR_!L!IbtS;jQXrwUK-1o7}gCQ8Vd^x)ZJV>K% zAntlmG*0mWC{em>)WI3$u^8qh_G{)Nj&XHFH$WrJKw&Jyp5=U@UW*D3XryJoOXTx% zG`JX~RXsnu`jg~h5RC15)YffplMOO+d-kv4M%4W<`hdn^|A*<~elfkxWrj9{BR zSH_lf@pCTlAXU#_agyG5fgzWQHD4l|3pCQB`r5JT)1-P$m)zaM8KstAb)I9d@bcEK z5#u$==PVk?VDdO#y()L?6%pyhYZ^0&AYt`kU?;CK@e$~+{^;AfIHL0xX@26&M2ay-NSi_)V=b3iPRNE#(6#h%6|(mSA$0>xTD z=&&)0{V9*Uix65par!u;Ojr!aDvun*`B~g7u|9t`rW>45V%;Pm1*<2n=s9%Hcd3Ar-)!Diugkw)8hTxPw-8tf$a5Ai**^m^Lv9nxrr%&^C$ zg0p4WPmu^QsO5O(bYKS|S8zNbITd;&j4XhPqsb-r_L=eqXrwqy!}0+;ZukWYDp#&k z&fD~GMyX2@plyH7NZdlFOlM9PXru)wj3g#aXv4@~am&5;w4B~W8Z9Q!5PYPNu@y1r zDh@e`s|aV_8>kUC!HJPHY3ea5M;V(mx;Q013E>84q+Mp%_Y#sgehzpR*T~Uloc4MH zG*Ug$95|0n0h^J0tAy5EVnrLo8T#*A?l)X@@J-v#Yv88m(%pY90JUhki>1!^uZ4i?Vt^BZZ$r zRxRO&u^fxV`G{CjIG_IpXrwqvf#d10zR=?C()d(uwN@8-&_;lP(m>I$j(GJrA}w3> za7L+%Ws`hN_@UIBjkC^$(*qhQCJ2(v*nkR5LU4~`aep4pC@~QI{N45poMn77%|zJzxGE(rA&tbScxtj3uvWNE8Npv_QR|vlJI-q#db(N}h(| z5PjMGi0x2XvI-w+#1NGRa;!09d)a!Hba9FP4(_0oSVhtakcHFGlR;?w!x5-1TlJAf zn~>BM83$w>YQSGoj>I}v$=*PXIAaDEIwWiEgq^USk5)=NvwAqA#A1D@!pz)O^@Wr0 z8kS7e1sZ9^46qTXq9@!28w{N(9hebf`@QU!5T%`}+u`E9ttWSZ2WisfVJiWs*i{_p z4aJn}vQrOF{f-=qi-ov26CY``@KAKj?~Yi?q{a%<&a5u*AdO+SIFfi1=0K8k*1k9>2WEr` zGg-!pLOr2G=Da&s2A>3a12j^g)QF|X!hwPmwz;*Q!kN><86{2-_Rz=+_on8-H;_YG z%<~p%#B~pgBxjLl2tRFU#B`4OKF%oHygtzn9Ik5U5!)+N^7)!v;6Y02%MiV^p%7ru zSZvg`cqbQVq?n(91hD~rCvhRqB}_U~Ixr)I)DvWy!vHow8^mzA1nY}yx`7%oG=tC` zLSG1Y@wzpSv}maVGeRsin>JW=?6T@1oU>9NXO!x++o8^7ijB=9gnT}@12e+Di->08 zro9^};ii-Ika>qRS`=squX%<+ZQf*RWv#&+m=QLaiIAw~*V>lhM6--i&l$)K@F3NyP6I!x zv5;n6Stw;%OGO7~lv?#2wcXKJlva1|xhx^mM;fhmN^!JuJhGGwKihN7Tuk?+H{t<}G+SqOR_1wP z8;Rq3*1~-OG}5y7BC_;9ru!8Z$9gw1r{@c&k2G2&J){z9d#fW!K2-wsv8J*f z&`1>p{eYZiUp`A!KF0D%z{44(HsZ9+ul81^O6O3{pF!_{MmqLyGI?fvUHf@$D~01= zxqiHP2|P$=n+VNJc|<;#T&G*T#91vaX_>K7e8mHfZVRiqx!NQ>#SW+Pdf zGOT^9&=j93{tKXyB2Uo>ljkEYw&h;ZlGHdbBZO?!?wjryK9?=IvX(4P-oqJX-`#?E z!^2Oo*QtqeT0gKjIHMeU(!LpUwN&r6+6JxNWd~-2 zySkMlA;{=#^g*#+_+FrlE?Yb}w$Rxn)YwLT7B!xsj1H}bK>L^xsPU^FDXd@99Na;f zf;KX>Mc${U)sAsX)Rt)J0*zGfEox6+zdW$R)=(&?D}5}z?_2N>{|6wGpn9V+003rT Bx4ZxV literal 0 HcmV?d00001 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( """\