diff --git a/tests/finance/test_slippage.py b/tests/finance/test_slippage.py index 2c6ae008..ba1ac919 100644 --- a/tests/finance/test_slippage.py +++ b/tests/finance/test_slippage.py @@ -27,20 +27,25 @@ from pandas.tslib import normalize_date from zipline.finance.slippage import VolumeShareSlippage -from zipline.protocol import DATASOURCE_TYPE +from zipline.protocol import DATASOURCE_TYPE, BarData from zipline.finance.blotter import Order - +from zipline.finance.asset_restrictions import NoRestrictions from zipline.data.data_portal import DataPortal -from zipline.protocol import BarData from zipline.testing import tmp_bcolz_equity_minute_bar_reader from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, WithSimParams, + WithTradingEnvironment, ZiplineTestCase, ) +from zipline.utils.classproperty import classproperty -class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): +class SlippageTestCase(WithCreateBarData, + 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 @@ -56,6 +61,10 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): freq='1min' ) + @classproperty + def CREATE_BARDATA_DATA_FREQUENCY(cls): + return cls.sim_params.data_frequency + @classmethod def make_equity_minute_bar_data(cls): yield 133, pd.DataFrame( @@ -74,97 +83,6 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): super(SlippageTestCase, cls).init_class_fixtures() cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133) - def test_volume_share_slippage(self): - 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_equity_minute_bar_reader(self.trading_calendar, days, assets) \ - as reader: - data_portal = DataPortal( - self.env.asset_finder, self.trading_calendar, - first_trading_day=reader.first_trading_day, - equity_minute_reader=reader, - ) - - slippage_model = VolumeShareSlippage() - - open_orders = [ - Order( - dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - amount=100, - filled=0, - sid=self.ASSET133 - ) - ] - - bar_data = BarData(data_portal, - lambda: self.minutes[0], - 'minute', - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 1) - _, txn = orders_txns[0] - - expected_txn = { - 'price': float(3.0001875), - 'dt': datetime.datetime( - 2006, 1, 5, 14, 31, tzinfo=pytz.utc), - 'amount': int(5), - 'sid': int(133), - 'commission': None, - 'type': DATASOURCE_TYPE.TRANSACTION, - 'order_id': open_orders[0].id - } - - self.assertIsNotNone(txn) - - # TODO: Make expected_txn an Transaction object and ensure there - # is a __eq__ for that class. - self.assertEquals(expected_txn, txn.__dict__) - - open_orders = [ - Order( - dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - amount=100, - filled=0, - sid=self.ASSET133 - ) - ] - - # Set bar_data to be a minute ahead of last trade. - # Volume share slippage should not execute when there is no trade. - bar_data = BarData(data_portal, - lambda: self.minutes[1], - 'minute', - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - def test_orders_limit(self): slippage_model = VolumeShareSlippage() slippage_model.data_portal = self.data_portal @@ -179,10 +97,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -202,10 +119,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -225,10 +141,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.6}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -265,10 +180,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -288,10 +202,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -311,10 +224,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.4}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -338,6 +250,376 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): for key, value in expected_txn.items(): self.assertEquals(value, txn[key]) + def test_orders_stop_limit(self): + slippage_model = VolumeShareSlippage() + slippage_model.data_portal = self.data_portal + + # long, does not trade + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': 100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 4.0, + 'limit': 3.0}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[2], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # long, does not trade - impacted price worse than limit price + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': 100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 4.0, + 'limit': 3.5}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[2], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # long, does trade + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': 100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 4.0, + 'limit': 3.6}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[2], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 1) + _, txn = orders_txns[0] + + expected_txn = { + 'price': float(3.50021875), + 'dt': datetime.datetime( + 2006, 1, 5, 14, 34, tzinfo=pytz.utc), + 'amount': int(50), + 'sid': int(133) + } + + for key, value in expected_txn.items(): + self.assertEquals(value, txn[key]) + + # short, does not trade + + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': -100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 3.0, + 'limit': 4.0}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # short, does not trade - impacted price worse than limit price + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': -100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 3.0, + 'limit': 3.5}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # short, does trade + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': -100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 3.0, + 'limit': 3.4}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 1) + _, txn = orders_txns[0] + + expected_txn = { + 'price': float(3.49978125), + 'dt': datetime.datetime( + 2006, 1, 5, 14, 32, tzinfo=pytz.utc), + 'amount': int(-50), + 'sid': int(133) + } + + for key, value in expected_txn.items(): + self.assertEquals(value, txn[key]) + + +class VolumeShareSlippageTestCase(WithCreateBarData, + 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' + ) + + @classproperty + def CREATE_BARDATA_DATA_FREQUENCY(cls): + return cls.sim_params.data_frequency + + @classmethod + def make_equity_minute_bar_data(cls): + yield 133, pd.DataFrame( + { + 'open': [3.00], + 'high': [3.15], + 'low': [2.85], + 'close': [3.00], + 'volume': [200], + }, + index=[cls.minutes[0]], + ) + + @classmethod + def init_class_fixtures(cls): + super(VolumeShareSlippageTestCase, cls).init_class_fixtures() + cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133) + + def test_volume_share_slippage(self): + + slippage_model = VolumeShareSlippage() + + open_orders = [ + Order( + dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + amount=100, + filled=0, + sid=self.ASSET133 + ) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 1) + _, txn = orders_txns[0] + + expected_txn = { + 'price': float(3.0001875), + 'dt': datetime.datetime( + 2006, 1, 5, 14, 31, tzinfo=pytz.utc), + 'amount': int(5), + 'sid': int(133), + 'commission': None, + 'type': DATASOURCE_TYPE.TRANSACTION, + 'order_id': open_orders[0].id + } + + self.assertIsNotNone(txn) + + # TODO: Make expected_txn an Transaction object and ensure there + # is a __eq__ for that class. + self.assertEquals(expected_txn, txn.__dict__) + + open_orders = [ + Order( + dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + amount=100, + filled=0, + sid=self.ASSET133 + ) + ] + + # Set bar_data to be a minute ahead of last trade. + # Volume share slippage should not execute when there is no trade. + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + +class OrdersStopTestCase(WithSimParams, + WithTradingEnvironment, + 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,) + minutes = pd.DatetimeIndex( + start=START_DATE, + end=END_DATE - pd.Timedelta('1 minute'), + freq='1min' + ) + + @classmethod + def init_class_fixtures(cls): + super(OrdersStopTestCase, cls).init_class_fixtures() + cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133) + STOP_ORDER_CASES = { # Stop orders can be long/short and have their price greater or # less than the stop. @@ -501,10 +783,14 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): try: dt = pd.Timestamp('2006-01-05 14:31', tz='UTC') - bar_data = BarData(data_portal, - lambda: dt, - 'minute', - self.trading_calendar) + bar_data = BarData( + data_portal, + lambda: dt, + self.sim_params.data_frequency, + self.trading_calendar, + NoRestrictions(), + ) + _, txn = next(slippage_model.simulate( bar_data, self.ASSET133, @@ -520,254 +806,3 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): for key, value in expected['transaction'].items(): self.assertEquals(value, txn[key]) - - def test_orders_stop_limit(self): - slippage_model = VolumeShareSlippage() - slippage_model.data_portal = self.data_portal - - # long, does not trade - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': 100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 4.0, - 'limit': 3.0}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[2], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # long, does not trade - impacted price worse than limit price - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': 100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 4.0, - 'limit': 3.5}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[2], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # long, does trade - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': 100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 4.0, - 'limit': 3.6}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[2], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 1) - _, txn = orders_txns[0] - - expected_txn = { - 'price': float(3.50021875), - 'dt': datetime.datetime( - 2006, 1, 5, 14, 34, tzinfo=pytz.utc), - 'amount': int(50), - 'sid': int(133) - } - - for key, value in expected_txn.items(): - self.assertEquals(value, txn[key]) - - # short, does not trade - - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': -100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 3.0, - 'limit': 4.0}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # short, does not trade - impacted price worse than limit price - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': -100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 3.0, - 'limit': 3.5}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # short, does trade - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': -100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 3.0, - 'limit': 3.4}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 1) - _, txn = orders_txns[0] - - expected_txn = { - 'price': float(3.49978125), - 'dt': datetime.datetime( - 2006, 1, 5, 14, 32, tzinfo=pytz.utc), - 'amount': int(-50), - 'sid': int(133) - } - - for key, value in expected_txn.items(): - self.assertEquals(value, txn[key]) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 22884a25..dea0ab45 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -76,6 +76,12 @@ from zipline.finance.commission import PerShare from zipline.finance.execution import LimitOrder from zipline.finance.order import ORDER_STATUS from zipline.finance.trading import SimulationParameters +from zipline.finance.asset_restrictions import ( + Restriction, + HistoricalRestrictions, + StaticRestrictions, + RESTRICTION_STATES, +) from zipline.testing import ( FakeDataPortal, create_daily_df_for_asset, @@ -122,6 +128,8 @@ from zipline.test_algorithms import ( SetMaxOrderCountAlgorithm, SetMaxOrderSizeAlgorithm, SetDoNotOrderListAlgorithm, + SetAssetRestrictionsAlgorithm, + SetMultipleAssetRestrictionsAlgorithm, SetMaxLeverageAlgorithm, api_algo, api_get_environment_algo, @@ -2788,34 +2796,114 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): env=self.env) self.check_algo_fails(algo, handle_data, 0) - def test_set_do_not_order_list(self): - # set the restricted list to be the sid, and fail. - algo = SetDoNotOrderListAlgorithm( - sid=self.sid, - restricted_list=[self.sid], - sim_params=self.sim_params, - env=self.env, - ) + def test_set_asset_restrictions(self): def handle_data(algo, data): + algo.could_trade = data.can_trade(algo.sid(self.sid)) algo.order(algo.sid(self.sid), 100) algo.order_count += 1 + # Set HistoricalRestrictions for one sid for the entire simulation, + # and fail. + rlm = HistoricalRestrictions([ + Restriction( + self.sid, + self.sim_params.start_session, + RESTRICTION_STATES.FROZEN) + ]) + algo = SetAssetRestrictionsAlgorithm( + sid=self.sid, + restrictions=rlm, + sim_params=self.sim_params, + env=self.env, + ) self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) + + # Set StaticRestrictions for one sid and fail. + rlm = StaticRestrictions([self.sid]) + algo = SetAssetRestrictionsAlgorithm( + sid=self.sid, + restrictions=rlm, + sim_params=self.sim_params, + env=self.env, + ) + self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) + + # just log an error on the violation if we choose not to fail. + algo = SetAssetRestrictionsAlgorithm( + sid=self.sid, + restrictions=rlm, + sim_params=self.sim_params, + env=self.env, + on_error='log' + ) + with make_test_handler(self) as log_catcher: + self.check_algo_succeeds(algo, handle_data) + logs = [r.message for r in log_catcher.records] + self.assertIn("Order for 100 shares of Equity(133 [A]) at " + "2006-01-03 21:00:00+00:00 violates trading constraint " + "RestrictedListOrder({})", logs) + self.assertFalse(algo.could_trade) # set the restricted list to exclude the sid, and succeed + rlm = HistoricalRestrictions([ + Restriction( + sid, + self.sim_params.start_session, + RESTRICTION_STATES.FROZEN) for sid in [134, 135, 136] + ]) + algo = SetAssetRestrictionsAlgorithm( + sid=self.sid, + restrictions=rlm, + sim_params=self.sim_params, + env=self.env, + ) + self.check_algo_succeeds(algo, handle_data) + self.assertTrue(algo.could_trade) + + @parameterized.expand([ + ('order_first_restricted_sid', 0), + ('order_second_restricted_sid', 1) + ]) + def test_set_multiple_asset_restrictions(self, name, to_order_idx): + + def handle_data(algo, data): + algo.could_trade1 = data.can_trade(algo.sid(self.sids[0])) + algo.could_trade2 = data.can_trade(algo.sid(self.sids[1])) + algo.order(algo.sid(self.sids[to_order_idx]), 100) + algo.order_count += 1 + + rl1 = StaticRestrictions([self.sids[0]]) + rl2 = StaticRestrictions([self.sids[1]]) + algo = SetMultipleAssetRestrictionsAlgorithm( + restrictions1=rl1, + restrictions2=rl2, + sim_params=self.sim_params, + env=self.env, + ) + self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade1) + self.assertFalse(algo.could_trade2) + + def test_set_do_not_order_list(self): + + def handle_data(algo, data): + algo.could_trade = data.can_trade(algo.sid(self.sid)) + algo.order(algo.sid(self.sid), 100) + algo.order_count += 1 + + rlm = [self.sid] algo = SetDoNotOrderListAlgorithm( sid=self.sid, - restricted_list=[134, 135, 136], + restricted_list=rlm, sim_params=self.sim_params, env=self.env, ) - def handle_data(algo, data): - algo.order(algo.sid(self.sid), 100) - algo.order_count += 1 - - self.check_algo_succeeds(algo, handle_data) + self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) def test_set_max_order_size(self): diff --git a/tests/test_api_shim.py b/tests/test_api_shim.py index 3aa64b8d..f79416d2 100644 --- a/tests/test_api_shim.py +++ b/tests/test_api_shim.py @@ -7,7 +7,6 @@ from pandas.core.common import PerformanceWarning from zipline import TradingAlgorithm from zipline.finance.trading import SimulationParameters -from zipline.protocol import BarData from zipline.testing import ( MockDailyBarReader, create_daily_df_for_asset, @@ -15,6 +14,7 @@ from zipline.testing import ( str_to_seconds, ) from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, WithSimParams, ZiplineTestCase, @@ -114,7 +114,11 @@ def handle_data(context, data): """ -class TestAPIShim(WithDataPortal, WithSimParams, ZiplineTestCase): +class TestAPIShim(WithCreateBarData, + 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' @@ -186,11 +190,8 @@ class TestAPIShim(WithDataPortal, WithSimParams, ZiplineTestCase): test_end_minute = self.trading_calendar.minutes_for_session( self.sim_params.sessions[0] )[-1] - bar_data = BarData( - self.data_portal, + bar_data = self.create_bardata( lambda: test_end_minute, - "minute", - self.trading_calendar ) ohlcvp_fields = [ "open", diff --git a/tests/test_bar_data.py b/tests/test_bar_data.py index f65a5456..e3ef7edc 100644 --- a/tests/test_bar_data.py +++ b/tests/test_bar_data.py @@ -23,8 +23,11 @@ import pandas as pd from zipline._protocol import handle_non_market_minutes -from zipline.data.data_portal import DataPortal -from zipline.protocol import BarData +from zipline.finance.asset_restrictions import ( + Restriction, + HistoricalRestrictions, + RESTRICTION_STATES, +) from zipline.testing import ( MockDailyBarReader, create_daily_df_for_asset, @@ -32,6 +35,7 @@ from zipline.testing import ( str_to_seconds, ) from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, ZiplineTestCase, ) @@ -49,6 +53,8 @@ field_info = { "close": 0 } +str_to_ts = lambda dt_str: pd.Timestamp(dt_str, tz='UTC') + class WithBarDataChecks(object): def assert_same(self, val1, val2): @@ -95,7 +101,8 @@ class WithBarDataChecks(object): getattr(bar_data, field) -class TestMinuteBarData(WithBarDataChecks, +class TestMinuteBarData(WithCreateBarData, + WithBarDataChecks, WithDataPortal, ZiplineTestCase): START_DATE = pd.Timestamp('2016-01-05', tz='UTC') @@ -205,8 +212,9 @@ class TestMinuteBarData(WithBarDataChecks, # this entire day is before either asset has started trading for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.check_internal_consistency(bar_data) self.assertFalse(bar_data.can_trade(self.ASSET1)) @@ -248,8 +256,9 @@ class TestMinuteBarData(WithBarDataChecks, # this test covers the "IPO morning" case, because asset2 only # has data starting on the 10th minute. - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.check_internal_consistency(bar_data) asset2_has_data = (((idx + 1) % 10) == 0) @@ -328,8 +337,9 @@ class TestMinuteBarData(WithBarDataChecks, # this is the last day the assets exist for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertTrue(bar_data.can_trade(self.ASSET1)) self.assertTrue(bar_data.can_trade(self.ASSET2)) @@ -347,8 +357,9 @@ class TestMinuteBarData(WithBarDataChecks, # this entire day is after both assets have stopped trading for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertFalse(bar_data.can_trade(self.ASSET1)) self.assertFalse(bar_data.can_trade(self.ASSET2)) @@ -390,8 +401,9 @@ class TestMinuteBarData(WithBarDataChecks, ) for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertEqual( idx + 1, bar_data.current(self.SPLIT_ASSET, "price") @@ -408,16 +420,16 @@ class TestMinuteBarData(WithBarDataChecks, ) for idx, minute in enumerate(day0_minutes[-10:-1]): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertEqual( 380, bar_data.current(self.ILLIQUID_SPLIT_ASSET, "price") ) - bar_data = BarData( - self.data_portal, lambda: day0_minutes[-1], "minute", - self.trading_calendar + bar_data = self.create_bardata( + lambda: day0_minutes[-1], ) self.assertEqual( @@ -426,8 +438,9 @@ class TestMinuteBarData(WithBarDataChecks, ) for idx, minute in enumerate(day1_minutes[0:9]): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) # should be half of 390, due to the split self.assertEqual( @@ -446,12 +459,12 @@ class TestMinuteBarData(WithBarDataChecks, tz='US/Eastern' ) - bar_data = BarData(self.data_portal, lambda: day, "minute", - self.trading_calendar) - bar_data2 = BarData(self.data_portal, - lambda: eight_fortyfive_am_eastern, - "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: day, + ) + bar_data2 = self.create_bardata( + lambda: eight_fortyfive_am_eastern, + ) with handle_non_market_minutes(bar_data), \ handle_non_market_minutes(bar_data2): @@ -482,20 +495,10 @@ class TestMinuteBarData(WithBarDataChecks, def test_get_value_during_non_market_hours(self): # make sure that if we try to get the OHLCV values of ASSET1 during # non-market hours, we don't get the previous market minute's values - futures_cal = get_calendar("us_futures") - data_portal = DataPortal( - self.env.asset_finder, - futures_cal, - first_trading_day=self.DATA_PORTAL_FIRST_TRADING_DAY, - equity_minute_reader=self.bcolz_equity_minute_bar_reader, - ) - - bar_data = BarData( - data_portal, - lambda: pd.Timestamp("2016-01-06 3:15", tz="US/Eastern"), - "minute", - futures_cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: + pd.Timestamp("2016-01-06 4:15", tz="US/Eastern"), ) self.assertTrue(np.isnan(bar_data.current(self.ASSET1, "open"))) @@ -508,14 +511,14 @@ class TestMinuteBarData(WithBarDataChecks, self.assertEqual(390, bar_data.current(self.ASSET1, "price")) def test_can_trade_equity_same_cal_outside_lifetime(self): - cal = get_calendar(self.ASSET1.exchange) # verify that can_trade returns False for the session before the # asset's first session - session_before_asset1_start = cal.previous_session_label( - self.ASSET1.start_date - ) - minutes_for_session = cal.minutes_for_session( + session_before_asset1_start = \ + self.trading_calendar.previous_session_label( + self.ASSET1.start_date + ) + minutes_for_session = self.trading_calendar.minutes_for_session( session_before_asset1_start ) @@ -526,14 +529,14 @@ class TestMinuteBarData(WithBarDataChecks, ) for minute in minutes_to_check: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertFalse(bar_data.can_trade(self.ASSET1)) # after asset lifetime - session_after_asset1_end = cal.next_session_label( + session_after_asset1_end = self.trading_calendar.next_session_label( self.ASSET1.end_date ) bts_after_asset1_end = session_after_asset1_end.replace( @@ -541,32 +544,32 @@ class TestMinuteBarData(WithBarDataChecks, ).tz_convert(None).tz_localize("US/Eastern") minutes_to_check = chain( - cal.minutes_for_session(session_after_asset1_end), + self.trading_calendar.minutes_for_session( + session_after_asset1_end + ), [bts_after_asset1_end] ) for minute in minutes_to_check: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertFalse(bar_data.can_trade(self.ASSET1)) def test_can_trade_equity_same_cal_exchange_closed(self): - cal = get_calendar(self.ASSET1.exchange) - # verify that can_trade returns true for minutes that are # outside the asset's calendar (assuming the asset is alive and # there is a last price), because the asset is alive on the # next market minute. - minutes = cal.minutes_for_sessions_in_range( + minutes = self.trading_calendar.minutes_for_sessions_in_range( self.ASSET1.start_date, self.ASSET1.end_date ) for minute in minutes: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertTrue(bar_data.can_trade(self.ASSET1)) @@ -576,13 +579,13 @@ class TestMinuteBarData(WithBarDataChecks, # 2016-01-05 15:20:00+00:00. Make sure that can_trade returns false # for all minutes in that session before the first trade, and true # for all minutes afterwards. - cal = get_calendar(self.ASSET1.exchange) - minutes_in_session = cal.minutes_for_session(self.ASSET1.start_date) + minutes_in_session = \ + self.trading_calendar.minutes_for_session(self.ASSET1.start_date) for minute in minutes_in_session[0:49]: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertFalse(bar_data.can_trade( @@ -590,14 +593,139 @@ class TestMinuteBarData(WithBarDataChecks, ) for minute in minutes_in_session[50:]: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertTrue(bar_data.can_trade( self.HILARIOUSLY_ILLIQUID_ASSET) ) + def test_is_stale_during_non_market_hours(self): + bar_data = self.create_bardata( + lambda: self.equity_minute_bar_days[1], + ) + + 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.adjustment_reader.get_adjustments_for_sid( + "splits", + self.SPLIT_ASSET.sid + ) + + self.assertEqual(1, len(splits)) + split = splits[0] + self.assertEqual( + split[0], + pd.Timestamp("2016-01-06", tz='UTC') + ) + + # Current day is 1/06/16 + day = self.equity_daily_bar_days[1] + eight_fortyfive_am_eastern = \ + pd.Timestamp("{0}-{1}-{2} 8:45".format( + day.year, day.month, day.day), + tz='US/Eastern' + ) + + bar_data = self.create_bardata( + lambda: eight_fortyfive_am_eastern, + ) + + expected = { + 'open': 391 / 2.0, + 'high': 392 / 2.0, + 'low': 389 / 2.0, + 'close': 390 / 2.0, + 'volume': 39000 * 2.0, + 'price': 390 / 2.0, + } + + with handle_non_market_minutes(bar_data): + for field in OHLCP + ['volume']: + value = bar_data.current(self.SPLIT_ASSET, field) + + # Assert the price is adjusted for the overnight split + self.assertEqual(value, expected[field]) + + def test_can_trade_restricted(self): + """ + Test that can_trade will return False for a sid if it is restricted + on that dt + """ + + minutes_to_check = [ + (str_to_ts("2016-01-05 14:31"), False), + (str_to_ts("2016-01-06 14:31"), False), + (str_to_ts("2016-01-07 14:31"), True), + (str_to_ts("2016-01-07 15:00"), False), + (str_to_ts("2016-01-07 15:30"), True), + ] + + rlm = HistoricalRestrictions([ + Restriction(1, str_to_ts('2016-01-05'), + RESTRICTION_STATES.FROZEN), + Restriction(1, str_to_ts('2016-01-07'), + RESTRICTION_STATES.ALLOWED), + Restriction(1, str_to_ts('2016-01-07 15:00'), + RESTRICTION_STATES.FROZEN), + Restriction(1, str_to_ts('2016-01-07 15:30'), + RESTRICTION_STATES.ALLOWED), + ]) + + for info in minutes_to_check: + bar_data = self.create_bardata( + simulation_dt_func=lambda: info[0], + restrictions=rlm, + ) + self.assertEqual(bar_data.can_trade(self.ASSET1), info[1]) + + +class TestMinuteBarDataMultipleExchanges(WithCreateBarData, + WithBarDataChecks, + 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] + + @classmethod + def make_equity_minute_bar_data(cls): + # asset1 has trades every minute + yield 1, create_minute_df_for_asset( + cls.trading_calendar, + cls.equity_minute_bar_days[0], + cls.equity_minute_bar_days[-1], + ) + + @classmethod + def make_futures_info(cls): + return pd.DataFrame.from_dict( + { + 6: { + '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'), + 'exchange': 'ICEUS', + }, + }, + orient='index', + ) + + @classmethod + def init_class_fixtures(cls): + super(TestMinuteBarDataMultipleExchanges, cls).init_class_fixtures() + cls.trading_calendar = get_calendar('CME') + def test_can_trade_multiple_exchange_closed(self): nyse_asset = self.asset_finder.retrieve_asset(1) ice_asset = self.asset_finder.retrieve_asset(6) @@ -639,70 +767,18 @@ class TestMinuteBarData(WithBarDataChecks, for info in minutes_to_check: # use the CME calendar, which covers 24 hours - bar_data = BarData(self.data_portal, lambda: info[0], "minute", - trading_calendar=get_calendar("CME")) + bar_data = self.create_bardata( + simulation_dt_func=lambda: info[0], + ) series = bar_data.can_trade([nyse_asset, ice_asset]) self.assertEqual(info[1], series.loc[nyse_asset]) self.assertEqual(info[2], series.loc[ice_asset]) - def test_is_stale_during_non_market_hours(self): - bar_data = BarData( - self.data_portal, - lambda: self.equity_minute_bar_days[1], - "minute", - self.trading_calendar - ) - 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.adjustment_reader.get_adjustments_for_sid( - "splits", - self.SPLIT_ASSET.sid - ) - - self.assertEqual(1, len(splits)) - split = splits[0] - self.assertEqual( - split[0], - pd.Timestamp("2016-01-06", tz='UTC') - ) - - # Current day is 1/06/16 - day = self.equity_daily_bar_days[1] - eight_fortyfive_am_eastern = \ - pd.Timestamp("{0}-{1}-{2} 8:45".format( - day.year, day.month, day.day), - tz='US/Eastern' - ) - - bar_data = BarData(self.data_portal, - lambda: eight_fortyfive_am_eastern, - "minute", - self.trading_calendar) - - expected = { - 'open': 391 / 2.0, - 'high': 392 / 2.0, - 'low': 389 / 2.0, - 'close': 390 / 2.0, - 'volume': 39000 * 2.0, - 'price': 390 / 2.0, - } - - with handle_non_market_minutes(bar_data): - for field in OHLCP + ['volume']: - value = bar_data.current(self.SPLIT_ASSET, field) - - # Assert the price is adjusted for the overnight split - self.assertEqual(value, expected[field]) - - -class TestDailyBarData(WithBarDataChecks, +class TestDailyBarData(WithCreateBarData, + WithBarDataChecks, WithDataPortal, ZiplineTestCase): START_DATE = pd.Timestamp('2016-01-05', tz='UTC') @@ -710,6 +786,7 @@ class TestDailyBarData(WithBarDataChecks, '2016-01-11', tz='UTC', ) + CREATE_BARDATA_DATA_FREQUENCY = 'daily' sids = ASSET_FINDER_EQUITY_SIDS = set(range(1, 9)) @@ -848,8 +925,9 @@ class TestDailyBarData(WithBarDataChecks, ) ) - bar_data = BarData(self.data_portal, lambda: minute, "daily", - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, + ) self.check_internal_consistency(bar_data) self.assertFalse(bar_data.can_trade(self.ASSET1)) @@ -871,13 +949,10 @@ class TestDailyBarData(WithBarDataChecks, def test_semi_active_day(self): # on self.equity_daily_bar_days[0], only asset1 has data - bar_data = BarData( - self.data_portal, - lambda: self.get_last_minute_of_session( + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.get_last_minute_of_session( self.equity_daily_bar_days[0] ), - "daily", - self.trading_calendar ) self.check_internal_consistency(bar_data) @@ -909,13 +984,10 @@ class TestDailyBarData(WithBarDataChecks, ) def test_fully_active_day(self): - bar_data = BarData( - self.data_portal, - lambda: self.get_last_minute_of_session( + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.get_last_minute_of_session( self.equity_daily_bar_days[1] ), - "daily", - self.trading_calendar ) self.check_internal_consistency(bar_data) @@ -936,13 +1008,10 @@ class TestDailyBarData(WithBarDataChecks, ) def test_last_active_day(self): - bar_data = BarData( - self.data_portal, - lambda: self.get_last_minute_of_session( + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.get_last_minute_of_session( self.equity_daily_bar_days[-1] ), - "daily", - self.trading_calendar ) self.check_internal_consistency(bar_data) @@ -971,8 +1040,9 @@ class TestDailyBarData(WithBarDataChecks, def test_after_assets_dead(self): session = self.END_DATE - bar_data = BarData(self.data_portal, lambda: session, "daily", - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: session, + ) self.check_internal_consistency(bar_data) for asset in self.ASSETS: @@ -1022,21 +1092,15 @@ class TestDailyBarData(WithBarDataChecks, ) # ... but that's it's not applied when using spot value - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[0], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[0], ) self.assertEqual( liquid_day_0_price, bar_data.current(liquid_asset, "price") ) - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[1], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[1], ) self.assertEqual( liquid_day_1_price, @@ -1045,21 +1109,15 @@ class TestDailyBarData(WithBarDataChecks, # ... 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.equity_daily_bar_days[1], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[1], ) self.assertEqual( illiquid_day_0_price, bar_data.current(illiquid_asset, "price") ) - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[2], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[2], ) # 3 (price from previous day) * 0.5 (split ratio) @@ -1067,3 +1125,29 @@ class TestDailyBarData(WithBarDataChecks, illiquid_day_1_price_adjusted, bar_data.current(illiquid_asset, "price") ) + + def test_can_trade_restricted(self): + """ + Test that can_trade will return False for a sid if it is restricted + on that dt + """ + + minutes_to_check = [ + (pd.Timestamp("2016-01-05", tz="UTC"), False), + (pd.Timestamp("2016-01-06", tz="UTC"), False), + (pd.Timestamp("2016-01-07", tz="UTC"), True), + ] + + rlm = HistoricalRestrictions([ + Restriction(1, str_to_ts('2016-01-05'), + RESTRICTION_STATES.FROZEN), + Restriction(1, str_to_ts('2016-01-07'), + RESTRICTION_STATES.ALLOWED), + ]) + + for info in minutes_to_check: + bar_data = self.create_bardata( + simulation_dt_func=lambda: info[0], + restrictions=rlm + ) + self.assertEqual(bar_data.can_trade(self.ASSET1), info[1]) diff --git a/tests/test_blotter.py b/tests/test_blotter.py index 89c251db..fa5d0cf3 100644 --- a/tests/test_blotter.py +++ b/tests/test_blotter.py @@ -31,8 +31,9 @@ from zipline.finance.slippage import ( DEFAULT_VOLUME_SLIPPAGE_BAR_LIMIT, FixedSlippage, ) -from zipline.protocol import BarData +from zipline.utils.classproperty import classproperty from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, WithLogger, WithSimParams, @@ -40,7 +41,8 @@ from zipline.testing.fixtures import ( ) -class BlotterTestCase(WithLogger, +class BlotterTestCase(WithCreateBarData, + WithLogger, WithDataPortal, WithSimParams, ZiplineTestCase): @@ -71,6 +73,10 @@ class BlotterTestCase(WithLogger, index=cls.sim_params.sessions, ) + @classproperty + def CREATE_BARDATA_DATA_FREQUENCY(cls): + return cls.sim_params.data_frequency + @parameterized.expand([(MarketOrder(), None, None), (LimitOrder(10), 10, None), (StopOrder(10), None, 10), @@ -219,11 +225,8 @@ class BlotterTestCase(WithLogger, filled_id = blotter.order(asset_24, 100, MarketOrder()) filled_order = None blotter.current_dt = self.sim_params.sessions[-1] - bar_data = BarData( - self.data_portal, - lambda: self.sim_params.sessions[-1], - self.sim_params.data_frequency, - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.sim_params.sessions[-1], ) txns, _, closed_orders = blotter.get_transactions(bar_data) for txn in txns: @@ -295,11 +298,8 @@ class BlotterTestCase(WithLogger, filled_order = None blotter.current_dt = dt - bar_data = BarData( - self.data_portal, - lambda: dt, - self.sim_params.data_frequency, - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: dt, ) txns, _, _ = blotter.get_transactions(bar_data) for txn in txns: diff --git a/tests/test_finance.py b/tests/test_finance.py index 4b374a35..bd84ba9b 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -37,6 +37,7 @@ 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.finance.asset_restrictions import NoRestrictions from zipline.protocol import BarData from zipline.testing import ( tmp_trading_env, @@ -317,10 +318,11 @@ class FinanceTestCase(WithLogger, order_date = order_date.replace(hour=14, minute=30) else: bar_data = BarData( - data_portal, - lambda: tick, - sim_params.data_frequency, - self.trading_calendar + data_portal=data_portal, + simulation_dt_func=lambda: tick, + data_frequency=sim_params.data_frequency, + trading_calendar=self.trading_calendar, + restrictions=NoRestrictions(), ) txns, _, closed_orders = blotter.get_transactions(bar_data) for txn in txns: diff --git a/tests/test_history.py b/tests/test_history.py index a68b7216..233d879d 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -21,20 +21,21 @@ import pandas as pd from six import iteritems from zipline import TradingAlgorithm -from zipline._protocol import handle_non_market_minutes +from zipline._protocol import handle_non_market_minutes, BarData from zipline.assets import Asset from zipline.errors import ( HistoryInInitialize, HistoryWindowStartsBeforeData, ) from zipline.finance.trading import SimulationParameters -from zipline.protocol import BarData +from zipline.finance.asset_restrictions import NoRestrictions from zipline.testing import ( create_minute_df_for_asset, str_to_seconds, MockDailyBarReader, ) from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, ZiplineTestCase, alias, @@ -46,7 +47,7 @@ OHLCP = OHLC + ['price'] ALL_FIELDS = OHLCP + ['volume'] -class WithHistory(WithDataPortal): +class WithHistory(WithCreateBarData, WithDataPortal): TRADING_START_DT = TRADING_ENV_MIN_DATE = START_DATE = pd.Timestamp( '2014-01-03', tz='UTC', @@ -251,8 +252,9 @@ class WithHistory(WithDataPortal): fields = fields if fields is not None else ALL_FIELDS assets = assets if assets is not None else [self.ASSET2, self.ASSET3] - bar_data = BarData(self.data_portal, lambda: dt, mode, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: dt, + ) check_internal_consistency( bar_data, assets, fields, 10, freq ) @@ -704,8 +706,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): )[0:60] for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, 'minute', - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) check_internal_consistency( bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, '1m' ) @@ -766,13 +769,12 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): ) )[1] - midnight_bar_data = \ - BarData(self.data_portal, lambda: midnight, 'minute', - self.trading_calendar) - - yesterday_bar_data = \ - BarData(self.data_portal, lambda: last_minute, 'minute', - self.trading_calendar) + midnight_bar_data = self.create_bardata( + lambda: midnight, + ) + yesterday_bar_data = self.create_bardata( + lambda: last_minute + ) with handle_non_market_minutes(midnight_bar_data): for field in ALL_FIELDS: @@ -789,8 +791,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): )[0:60] for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, 'minute', - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute + ) check_internal_consistency( bar_data, self.SHORT_ASSET, ALL_FIELDS, 30, '1m' ) @@ -799,8 +802,13 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): data_portal = self.make_data_portal() # choose a window that contains the last minute of the asset - bar_data = BarData(data_portal, lambda: minutes[15], 'minute', - self.trading_calendar) + bar_data = BarData( + data_portal=data_portal, + simulation_dt_func=lambda: minutes[15], + data_frequency='minute', + restrictions=NoRestrictions(), + trading_calendar=self.trading_calendar, + ) # close high low open price volume # 2015-01-06 20:47:00+00:00 768 770 767 769 768 76800 @@ -1012,8 +1020,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): def test_passing_iterable_to_history_regular_hours(self): # regular hours current_dt = pd.Timestamp("2015-01-06 9:45", tz='US/Eastern') - bar_data = BarData(self.data_portal, lambda: current_dt, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: current_dt, + ) bar_data.history(pd.Index([self.ASSET1, self.ASSET2]), "high", 5, "1m") @@ -1021,8 +1030,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): def test_passing_iterable_to_history_bts(self): # before market hours current_dt = pd.Timestamp("2015-01-07 8:45", tz='US/Eastern') - bar_data = BarData(self.data_portal, lambda: current_dt, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: current_dt, + ) with handle_non_market_minutes(bar_data): bar_data.history(pd.Index([self.ASSET1, self.ASSET2]), @@ -1031,8 +1041,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): 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', - self.trading_calendar) + bar_data = self.create_bardata( + lambda: current_dt, + ) adj_expected = { 'open': np.arange(8381, 8391) / 4.0, @@ -1341,6 +1352,8 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): + CREATE_BARDATA_DATA_FREQUENCY = 'daily' + @classmethod def make_equity_daily_bar_data(cls): yield 1, cls.create_df_for_asset( @@ -1403,8 +1416,9 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): ) for idx, day in enumerate(days): - bar_data = BarData(self.data_portal, lambda: day, 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: day, + ) check_internal_consistency( bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, '1d' ) @@ -1445,10 +1459,9 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): # 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', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: pd.Timestamp('2016-01-06', tz='UTC'), + ) for field in OHLCP: window = bar_data.history( @@ -1486,8 +1499,9 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): # days has 1/7, 1/8 for idx, day in enumerate(days): - bar_data = BarData(self.data_portal, lambda: day, 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: day, + ) check_internal_consistency( bar_data, self.SHORT_ASSET, ALL_FIELDS, 2, '1d' ) @@ -1639,10 +1653,10 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): # asset1 ends on 2016-01-30 # asset2 ends on 2016-01-04 - bar_data = BarData(self.data_portal, - lambda: pd.Timestamp('2016-01-06 16:00', tz='UTC'), - 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: + pd.Timestamp('2016-01-06 16:00', tz='UTC'), + ) for field in OHLCP: window = bar_data.history( diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py new file mode 100644 index 00000000..9ec90e66 --- /dev/null +++ b/tests/test_restrictions.py @@ -0,0 +1,422 @@ +import pandas as pd +from pandas.util.testing import assert_series_equal +from six import iteritems +from functools import partial + +from toolz import groupby + +from zipline.finance.asset_restrictions import ( + RESTRICTION_STATES, + Restriction, + HistoricalRestrictions, + StaticRestrictions, + SecurityListRestrictions, + NoRestrictions, + _UnionRestrictions, +) + +from zipline.testing import parameter_space +from zipline.testing.fixtures import ( + WithDataPortal, + ZiplineTestCase, +) + +str_to_ts = lambda dt_str: pd.Timestamp(dt_str, tz='UTC') +FROZEN = RESTRICTION_STATES.FROZEN +ALLOWED = RESTRICTION_STATES.ALLOWED +MINUTE = pd.Timedelta(minutes=1) + + +class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): + + ASSET_FINDER_EQUITY_SIDS = 1, 2, 3 + + @classmethod + def init_class_fixtures(cls): + super(RestrictionsTestCase, cls).init_class_fixtures() + 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.ALL_ASSETS = [cls.ASSET1, cls.ASSET2, cls.ASSET3] + + def assert_is_restricted(self, rl, asset, dt): + self.assertTrue(rl.is_restricted(asset, dt)) + + def assert_not_restricted(self, rl, asset, dt): + self.assertFalse(rl.is_restricted(asset, dt)) + + def assert_all_restrictions(self, rl, expected, dt): + self.assert_many_restrictions(rl, self.ALL_ASSETS, expected, dt) + + def assert_many_restrictions(self, rl, assets, expected, dt): + assert_series_equal( + rl.is_restricted(assets, dt), + pd.Series(index=pd.Index(assets), data=expected), + ) + + @parameter_space( + date_offset=( + pd.Timedelta(0), + pd.Timedelta('1 minute'), + pd.Timedelta('15 hours 5 minutes') + ), + check_unordered=(False, True), + __fail_fast=True, + ) + def test_historical_restrictions(self, date_offset, check_unordered): + """ + Test historical restrictions for both interday and intraday + restrictions, as well as restrictions defined in/not in order, for both + single- and multi-asset queries + """ + if check_unordered: + def maybe_scramble(rs): + # Swap the first two restrictions to check that we don't care + # that the restriction dates are ordered. + tmp = rs[0] + rs[0] = rs[1] + rs[1] = tmp + return rs + else: + maybe_scramble = lambda r: r + + def rdate(s): + """Convert a date string into a restriction for that date.""" + # Add date_offset to check that we handle intraday changes. + return str_to_ts(s) + date_offset + + all_restrictions = ( + maybe_scramble([ + Restriction(self.ASSET1, rdate('2011-01-04'), FROZEN), + Restriction(self.ASSET1, rdate('2011-01-05'), ALLOWED), + Restriction(self.ASSET1, rdate('2011-01-06'), FROZEN), + ]) + + + maybe_scramble([ + Restriction(self.ASSET2, rdate('2011-01-05'), FROZEN), + Restriction(self.ASSET2, rdate('2011-01-06'), ALLOWED), + Restriction(self.ASSET2, rdate('2011-01-07'), FROZEN), + ]) + ) + restrictions_by_asset = groupby(lambda r: r.asset, all_restrictions) + + rl = HistoricalRestrictions(all_restrictions) + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + assert_all_restrictions = partial(self.assert_all_restrictions, rl) + + # Check individual restrictions. + for asset, r_history in iteritems(restrictions_by_asset): + freeze_dt, unfreeze_dt, re_freeze_dt = ( + sorted([r.effective_date for r in r_history]) + ) + + # Starts implicitly unrestricted. Restricted on or after the freeze + assert_not_restricted(asset, freeze_dt - MINUTE) + assert_is_restricted(asset, freeze_dt) + assert_is_restricted(asset, freeze_dt + MINUTE) + + # Unrestricted on or after the unfreeze + assert_is_restricted(asset, unfreeze_dt - MINUTE) + assert_not_restricted(asset, unfreeze_dt) + assert_not_restricted(asset, unfreeze_dt + MINUTE) + + # Restricted again on or after the freeze + assert_not_restricted(asset, re_freeze_dt - MINUTE) + assert_is_restricted(asset, re_freeze_dt) + assert_is_restricted(asset, re_freeze_dt + MINUTE) + + # Should stay restricted for the rest of time + assert_is_restricted(asset, re_freeze_dt + MINUTE * 1000000) + + # Check vectorized restrictions. + # Expected results for [self.ASSET1, self.ASSET2, self.ASSET3], + # ASSET3 is always False as it has no defined restrictions + + # 01/04 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: ALLOWED + d0 = rdate('2011-01-04') + assert_all_restrictions([False, False, False], d0 - MINUTE) + assert_all_restrictions([True, False, False], d0) + assert_all_restrictions([True, False, False], d0 + MINUTE) + + # 01/05 XX:00 ASSET1: FROZEN --> ALLOWED; ASSET2: ALLOWED --> FROZEN + d1 = rdate('2011-01-05') + assert_all_restrictions([True, False, False], d1 - MINUTE) + assert_all_restrictions([False, True, False], d1) + assert_all_restrictions([False, True, False], d1 + MINUTE) + + # 01/06 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: FROZEN --> ALLOWED + d2 = rdate('2011-01-06') + assert_all_restrictions([False, True, False], d2 - MINUTE) + assert_all_restrictions([True, False, False], d2) + assert_all_restrictions([True, False, False], d2 + MINUTE) + + # 01/07 XX:00 ASSET1: FROZEN; ASSET2: ALLOWED --> FROZEN + d3 = rdate('2011-01-07') + assert_all_restrictions([True, False, False], d3 - MINUTE) + assert_all_restrictions([True, True, False], d3) + assert_all_restrictions([True, True, False], d3 + MINUTE) + + # Should stay restricted for the rest of time + assert_all_restrictions( + [True, True, False], + d3 + (MINUTE * 10000000) + ) + + def test_historical_restrictions_consecutive_states(self): + """ + Test that defining redundant consecutive restrictions still works + """ + rl = HistoricalRestrictions([ + Restriction(self.ASSET1, str_to_ts('2011-01-04'), ALLOWED), + Restriction(self.ASSET1, str_to_ts('2011-01-05'), ALLOWED), + Restriction(self.ASSET1, str_to_ts('2011-01-06'), FROZEN), + Restriction(self.ASSET1, str_to_ts('2011-01-07'), FROZEN), + ]) + + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + + # (implicit) ALLOWED --> ALLOWED + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04') - MINUTE) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04')) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04') + MINUTE) + + # ALLOWED --> ALLOWED + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05') - MINUTE) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05')) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05') + MINUTE) + + # ALLOWED --> FROZEN + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-06') - MINUTE) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-06')) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-06') + MINUTE) + + # FROZEN --> FROZEN + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07') - MINUTE) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07')) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07') + MINUTE) + + def test_static_restrictions(self): + """ + Test single- and multi-asset queries on static restrictions + """ + + restricted_a1 = self.ASSET1 + restricted_a2 = self.ASSET2 + unrestricted_a3 = self.ASSET3 + + rl = StaticRestrictions([restricted_a1, restricted_a2]) + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + assert_all_restrictions = partial(self.assert_all_restrictions, rl) + + for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03', + '2011-01-04', + '2011-01-04 1:01', + '2020-01-04')]: + assert_is_restricted(restricted_a1, dt) + assert_is_restricted(restricted_a2, dt) + assert_not_restricted(unrestricted_a3, dt) + + assert_all_restrictions([True, True, False], dt) + + def test_security_list_restrictions(self): + """ + Test single- and multi-asset queries on restrictions defined by + zipline.utils.security_list.SecurityList + """ + + # A mock SecurityList object filled with fake data + class SecurityList(object): + def __init__(self, assets_by_dt): + self.assets_by_dt = assets_by_dt + + def current_securities(self, dt): + return self.assets_by_dt[dt] + + assets_by_dt = { + str_to_ts('2011-01-03'): [self.ASSET1], + str_to_ts('2011-01-04'): [self.ASSET2, self.ASSET3], + str_to_ts('2011-01-05'): [self.ASSET1, self.ASSET2, self.ASSET3], + } + + rl = SecurityListRestrictions(SecurityList(assets_by_dt)) + + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + assert_all_restrictions = partial(self.assert_all_restrictions, rl) + + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-03')) + assert_not_restricted(self.ASSET2, str_to_ts('2011-01-03')) + assert_not_restricted(self.ASSET3, str_to_ts('2011-01-03')) + assert_all_restrictions( + [True, False, False], str_to_ts('2011-01-03') + ) + + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04')) + assert_is_restricted(self.ASSET2, str_to_ts('2011-01-04')) + assert_is_restricted(self.ASSET3, str_to_ts('2011-01-04')) + assert_all_restrictions( + [False, True, True], str_to_ts('2011-01-04') + ) + + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-05')) + assert_is_restricted(self.ASSET2, str_to_ts('2011-01-05')) + assert_is_restricted(self.ASSET3, str_to_ts('2011-01-05')) + assert_all_restrictions( + [True, True, True], + str_to_ts('2011-01-05') + ) + + def test_noop_restrictions(self): + """ + Test single- and multi-asset queries on no-op restrictions + """ + + rl = NoRestrictions() + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_all_restrictions = partial(self.assert_all_restrictions, rl) + + for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03', + '2011-01-04', + '2020-01-04')]: + assert_not_restricted(self.ASSET1, dt) + assert_not_restricted(self.ASSET2, dt) + assert_not_restricted(self.ASSET3, dt) + assert_all_restrictions([False, False, False], dt) + + def test_union_restrictions(self): + """ + Test that we appropriately union restrictions together, including + eliminating redundancy (ignoring NoRestrictions) and flattening out + the underlying sub-restrictions of _UnionRestrictions + """ + + no_restrictions_rl = NoRestrictions() + + st_restrict_asset1 = StaticRestrictions([self.ASSET1]) + st_restrict_asset2 = StaticRestrictions([self.ASSET2]) + st_restricted_assets = [self.ASSET1, self.ASSET2] + + before_frozen_dt = str_to_ts('2011-01-05') + freeze_dt_1 = str_to_ts('2011-01-06') + unfreeze_dt = str_to_ts('2011-01-06 16:00') + hist_restrict_asset3_1 = HistoricalRestrictions([ + Restriction(self.ASSET3, freeze_dt_1, FROZEN), + Restriction(self.ASSET3, unfreeze_dt, ALLOWED) + ]) + + freeze_dt_2 = str_to_ts('2011-01-07') + hist_restrict_asset3_2 = HistoricalRestrictions([ + Restriction(self.ASSET3, freeze_dt_2, FROZEN) + ]) + + # A union of a NoRestrictions with a non-trivial restriction should + # yield the original restriction + trivial_union_restrictions = no_restrictions_rl | st_restrict_asset1 + self.assertIsInstance(trivial_union_restrictions, StaticRestrictions) + + # A union of two non-trivial restrictions should yield a + # UnionRestrictions + st_union_restrictions = st_restrict_asset1 | st_restrict_asset2 + self.assertIsInstance(st_union_restrictions, _UnionRestrictions) + + arb_dt = str_to_ts('2011-01-04') + self.assert_is_restricted(st_restrict_asset1, self.ASSET1, arb_dt) + self.assert_not_restricted(st_restrict_asset1, self.ASSET2, arb_dt) + self.assert_not_restricted(st_restrict_asset2, self.ASSET1, arb_dt) + self.assert_is_restricted(st_restrict_asset2, self.ASSET2, arb_dt) + self.assert_is_restricted(st_union_restrictions, self.ASSET1, arb_dt) + self.assert_is_restricted(st_union_restrictions, self.ASSET2, arb_dt) + self.assert_many_restrictions( + st_restrict_asset1, + st_restricted_assets, + [True, False], + arb_dt + ) + self.assert_many_restrictions( + st_restrict_asset2, + st_restricted_assets, + [False, True], + arb_dt + ) + self.assert_many_restrictions( + st_union_restrictions, + st_restricted_assets, + [True, True], + arb_dt + ) + + # A union of a 2-sub-restriction UnionRestrictions and a + # non-trivial restrictions should yield a UnionRestrictions with + # 3 sub restrictions. Works with UnionRestrictions on both the left + # side or right side + for r1, r2 in [ + (st_union_restrictions, hist_restrict_asset3_1), + (hist_restrict_asset3_1, st_union_restrictions) + ]: + union_or_hist_restrictions = r1 | r2 + self.assertIsInstance( + union_or_hist_restrictions, _UnionRestrictions) + self.assertEqual( + len(union_or_hist_restrictions.sub_restrictions), 3) + + # Includes the two static restrictions on ASSET1 and ASSET2, + # and the historical restriction on ASSET3 starting on freeze_dt_1 + # and ending on unfreeze_dt + self.assert_all_restrictions( + union_or_hist_restrictions, + [True, True, False], + before_frozen_dt + ) + self.assert_all_restrictions( + union_or_hist_restrictions, + [True, True, True], + freeze_dt_1 + ) + self.assert_all_restrictions( + union_or_hist_restrictions, + [True, True, False], + unfreeze_dt + ) + self.assert_all_restrictions( + union_or_hist_restrictions, + [True, True, False], + freeze_dt_2 + ) + + # A union of two 2-sub-restrictions UnionRestrictions should yield a + # UnionRestrictions with 4 sub restrictions. + hist_union_restrictions = \ + hist_restrict_asset3_1 | hist_restrict_asset3_2 + multi_union_restrictions = \ + st_union_restrictions | hist_union_restrictions + + self.assertIsInstance(multi_union_restrictions, _UnionRestrictions) + self.assertEqual(len(multi_union_restrictions.sub_restrictions), 4) + + # Includes the two static restrictions on ASSET1 and ASSET2, the + # first historical restriction on ASSET3 starting on freeze_dt_1 and + # ending on unfreeze_dt, and the second historical restriction on + # ASSET3 starting on freeze_dt_2 + self.assert_all_restrictions( + multi_union_restrictions, + [True, True, False], + before_frozen_dt + ) + self.assert_all_restrictions( + multi_union_restrictions, + [True, True, True], + freeze_dt_1 + ) + self.assert_all_restrictions( + multi_union_restrictions, + [True, True, False], + unfreeze_dt + ) + self.assert_all_restrictions( + multi_union_restrictions, + [True, True, True], + freeze_dt_2 + ) diff --git a/tests/test_security_list.py b/tests/test_security_list.py index f0609142..b381adc3 100644 --- a/tests/test_security_list.py +++ b/tests/test_security_list.py @@ -2,6 +2,7 @@ from datetime import timedelta import pandas as pd from testfixtures import TempDirectory +from nose_parameterized import parameterized from zipline.algorithm import TradingAlgorithm from zipline.errors import TradingControlViolation @@ -29,19 +30,32 @@ LEVERAGED_ETFS = load_from_directory('leveraged_etf_list') class RestrictedAlgoWithCheck(TradingAlgorithm): def initialize(self, symbol): self.rl = SecurityListSet(self.get_datetime, self.asset_finder) - self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.set_asset_restrictions(self.rl.restrict_leveraged_etfs) self.order_count = 0 self.sid = self.symbol(symbol) def handle_data(self, data): if not self.order_count: if self.sid not in \ - self.rl.leveraged_etf_list: + self.rl.leveraged_etf_list.\ + current_securities(self.get_datetime()): self.order(self.sid, 100) self.order_count += 1 class RestrictedAlgoWithoutCheck(TradingAlgorithm): + def initialize(self, symbol): + self.rl = SecurityListSet(self.get_datetime, self.asset_finder) + self.set_asset_restrictions(self.rl.restrict_leveraged_etfs) + self.order_count = 0 + self.sid = self.symbol(symbol) + + def handle_data(self, data): + self.order(self.sid, 100) + self.order_count += 1 + + +class RestrictedAlgoWithoutCheckSetDoNotOrderList(TradingAlgorithm): def initialize(self, symbol): self.rl = SecurityListSet(self.get_datetime, self.asset_finder) self.set_do_not_order_list(self.rl.leveraged_etf_list) @@ -56,13 +70,14 @@ class RestrictedAlgoWithoutCheck(TradingAlgorithm): class IterateRLAlgo(TradingAlgorithm): def initialize(self, symbol): self.rl = SecurityListSet(self.get_datetime, self.asset_finder) - self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.set_asset_restrictions(self.rl.restrict_leveraged_etfs) self.order_count = 0 self.sid = self.symbol(symbol) self.found = False def handle_data(self, data): - for stock in self.rl.leveraged_etf_list: + for stock in self.rl.leveraged_etf_list.\ + current_securities(self.get_datetime()): if stock == self.sid: self.found = True @@ -151,7 +166,8 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): for symbol in ["BZQ", "URTY", "JFT"]] ] for sid in should_exist: - self.assertIn(sid, rl.leveraged_etf_list) + self.assertIn( + sid, rl.leveraged_etf_list.current_securities(get_datetime())) # assert that a sample of allowed stocks are not in restricted shouldnt_exist = [ @@ -162,7 +178,8 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): for symbol in ["AAPL", "GOOG"]] ] for sid in shouldnt_exist: - self.assertNotIn(sid, rl.leveraged_etf_list) + self.assertNotIn( + sid, rl.leveraged_etf_list.current_securities(get_datetime())) def test_security_add(self): def get_datetime(): @@ -178,15 +195,24 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): ) for symbol in ["AAPL", "GOOG", "BZQ", "URTY"]] ] for sid in should_exist: - self.assertIn(sid, rl.leveraged_etf_list) + self.assertIn( + sid, + rl.leveraged_etf_list.current_securities(get_datetime()) + ) def test_security_add_delete(self): with security_list_copy(): def get_datetime(): return pd.Timestamp("2015-01-27", tz='UTC') rl = SecurityListSet(get_datetime, self.env.asset_finder) - self.assertNotIn("BZQ", rl.leveraged_etf_list) - self.assertNotIn("URTY", rl.leveraged_etf_list) + self.assertNotIn( + "BZQ", + rl.leveraged_etf_list.current_securities(get_datetime()) + ) + self.assertNotIn( + "URTY", + rl.leveraged_etf_list.current_securities(get_datetime()) + ) def test_algo_without_rl_violation_via_check(self): algo = RestrictedAlgoWithCheck(symbol='BZQ', @@ -200,10 +226,15 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): env=self.env) algo.run(self.data_portal) - def test_algo_with_rl_violation(self): - algo = RestrictedAlgoWithoutCheck(symbol='BZQ', - sim_params=self.sim_params, - env=self.env) + @parameterized.expand([ + ('using_set_do_not_order_list', + RestrictedAlgoWithoutCheckSetDoNotOrderList), + ('using_set_restrictions', RestrictedAlgoWithoutCheck), + ]) + def test_algo_with_rl_violation(self, name, algo_class): + algo = algo_class(symbol='BZQ', + sim_params=self.sim_params, + env=self.env) with self.assertRaises(TradingControlViolation) as ctx: algo.run(self.data_portal) diff --git a/tests/test_tradesimulation.py b/tests/test_tradesimulation.py index d3bb8b33..6277b8e2 100644 --- a/tests/test_tradesimulation.py +++ b/tests/test_tradesimulation.py @@ -24,6 +24,7 @@ from zipline import TradingAlgorithm from zipline.gens.sim_engine import BEFORE_TRADING_START_BAR from zipline.finance.performance import PerformanceTracker +from zipline.finance.asset_restrictions import NoRestrictions from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.sources.benchmark_source import BenchmarkSource from zipline.test_algorithms import NoopAlgorithm @@ -135,6 +136,7 @@ def initialize(context): self.data_portal, BeforeTradingStartsOnlyClock(dt), algo._create_benchmark_source(), + NoRestrictions(), None ) diff --git a/zipline/_protocol.pyx b/zipline/_protocol.pyx index fd0c39b2..43e7d20d 100644 --- a/zipline/_protocol.pyx +++ b/zipline/_protocol.pyx @@ -153,6 +153,9 @@ cdef class BarData: data_frequency : {'minute', 'daily'} The frequency of the bar data; i.e. whether the data is daily or minute bars + restrictions : zipline.finance.asset_restrictions.Restrictions + Object that combines and returns restricted list information from + multiple sources universe_func : callable, optional Function which returns the current 'universe'. This is for backwards compatibility with older API concepts. @@ -160,17 +163,19 @@ cdef class BarData: cdef object data_portal cdef object simulation_dt_func cdef object data_frequency + cdef object restrictions cdef dict _views cdef object _universe_func cdef object _last_calculated_universe cdef object _universe_last_updated_at cdef bool _daily_mode cdef object _trading_calendar + cdef object _is_restricted cdef bool _adjust_minutes def __init__(self, data_portal, simulation_dt_func, data_frequency, - trading_calendar, universe_func=None): + trading_calendar, restrictions, universe_func=None): self.data_portal = data_portal self.simulation_dt_func = simulation_dt_func self.data_frequency = data_frequency @@ -185,6 +190,7 @@ cdef class BarData: self._adjust_minutes = False self._trading_calendar = trading_calendar + self._is_restricted = restrictions.is_restricted cdef _get_equity_price_view(self, asset): """ @@ -482,6 +488,9 @@ cdef class BarData: cdef object session_label cdef object dt_to_use_for_exchange_check, + if self._is_restricted(asset, adjusted_dt): + return False + session_label = self._trading_calendar.minute_to_session_label(dt) if not asset.is_alive_for_session(session_label): diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 86553d27..dbfb373a 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -76,11 +76,17 @@ from zipline.finance.execution import ( StopOrder, ) from zipline.finance.performance import PerformanceTracker +from zipline.finance.asset_restrictions import Restrictions from zipline.finance.slippage import ( VolumeShareSlippage, SlippageModel ) from zipline.finance.cancel_policy import NeverCancel, CancelPolicy +from zipline.finance.asset_restrictions import ( + NoRestrictions, + StaticRestrictions, + SecurityListRestrictions, +) from zipline.assets import Asset, Future from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.pipeline import Pipeline @@ -120,6 +126,7 @@ from zipline.utils.math_utils import ( round_if_near_integer ) from zipline.utils.preprocess import preprocess +from zipline.utils.security_list import SecurityList import zipline.protocol from zipline.sources.requests_csv import PandasRequestsCSV @@ -418,6 +425,8 @@ class TradingAlgorithm(object): # A dictionary of the actual capital change deltas, keyed by timestamp self.capital_change_deltas = {} + self.restrictions = NoRestrictions() + def init_engine(self, get_loader): """ Construct and store a PipelineEngine from loader. @@ -564,6 +573,7 @@ class TradingAlgorithm(object): self.data_portal, self._create_clock(), self._create_benchmark_source(), + self.restrictions, universe_func=self._calculate_universe ) @@ -2083,7 +2093,8 @@ class TradingAlgorithm(object): def set_max_position_size(self, asset=None, max_shares=None, - max_notional=None): + max_notional=None, + on_error='fail'): """Set a limit on the number of shares and/or dollar value held for the given sid. Limits are treated as absolute values and are enforced at the time that the algo attempts to place an order for sid. This means @@ -2107,14 +2118,16 @@ class TradingAlgorithm(object): """ control = MaxPositionSize(asset=asset, max_shares=max_shares, - max_notional=max_notional) + max_notional=max_notional, + on_error=on_error) self.register_trading_control(control) @api_method def set_max_order_size(self, asset=None, max_shares=None, - max_notional=None): + max_notional=None, + on_error='fail'): """Set a limit on the number of shares and/or dollar value of any single order placed for sid. Limits are treated as absolute values and are enforced at the time that the algo attempts to place an order for sid. @@ -2134,11 +2147,12 @@ class TradingAlgorithm(object): """ control = MaxOrderSize(asset=asset, max_shares=max_shares, - max_notional=max_notional) + max_notional=max_notional, + on_error=on_error) self.register_trading_control(control) @api_method - def set_max_order_count(self, max_count): + def set_max_order_count(self, max_count, on_error='fail'): """Set a limit on the number of orders that can be placed in a single day. @@ -2147,27 +2161,68 @@ class TradingAlgorithm(object): max_count : int The maximum number of orders that can be placed on any single day. """ - control = MaxOrderCount(max_count) + control = MaxOrderCount(on_error, max_count) self.register_trading_control(control) @api_method - def set_do_not_order_list(self, restricted_list): + def set_do_not_order_list(self, restricted_list, on_error='fail'): """Set a restriction on which assets can be ordered. Parameters ---------- - restricted_list : container[Asset] + restricted_list : container[Asset], SecurityList The assets that cannot be ordered. """ - control = RestrictedListOrder(restricted_list) - self.register_trading_control(control) + if isinstance(restricted_list, SecurityList): + warnings.warn( + "`set_do_not_order_list(security_lists.leveraged_etf_list)` " + "is deprecated. Use `set_asset_restrictions(" + "security_lists.restrict_leveraged_etfs)` instead.", + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + restrictions = SecurityListRestrictions(restricted_list) + else: + warnings.warn( + "`set_do_not_order_list(container_of_assets)` is deprecated. " + "Create a zipline.finance.asset_restrictions." + "StaticRestrictions object with a container of assets and use " + "`set_asset_restrictions(StaticRestrictions(" + "container_of_assets))` instead.", + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + restrictions = StaticRestrictions(restricted_list) + + self.set_asset_restrictions(restrictions, on_error) @api_method - def set_long_only(self): + @expect_types( + restrictions=Restrictions, + on_error=str, + ) + def set_asset_restrictions(self, restrictions, on_error='fail'): + """Set a restriction on which assets can be ordered. + + Parameters + ---------- + restricted_list : Restrictions + An object providing information about restricted assets. + + See Also + -------- + zipline.finance.asset_restrictions.Restrictions + """ + control = RestrictedListOrder(on_error, restrictions) + self.register_trading_control(control) + self.restrictions |= restrictions + + @api_method + def set_long_only(self, on_error='fail'): """Set a rule specifying that this algorithm cannot take short positions. """ - self.register_trading_control(LongOnly()) + self.register_trading_control(LongOnly(on_error)) ############## # Pipeline API diff --git a/zipline/api.py b/zipline/api.py index ef3a75f8..68b78714 100644 --- a/zipline/api.py +++ b/zipline/api.py @@ -16,6 +16,12 @@ # Note that part of the API is implemented in TradingAlgorithm as # methods (e.g. order). These are added to this namespace via the # decorator ``api_method`` inside of algorithm.py. +from .finance.asset_restrictions import ( + Restriction, + StaticRestrictions, + HistoricalRestrictions, + RESTRICTION_STATES, +) from .finance import commission, execution, slippage, cancel_policy from .finance.cancel_policy import ( NeverCancel, @@ -36,6 +42,10 @@ __all__ = [ 'FixedSlippage', 'NeverCancel', 'VolumeShareSlippage', + 'Restriction', + 'StaticRestrictions', + 'HistoricalRestrictions', + 'RESTRICTION_STATES', 'cancel_policy', 'commission', 'date_rules', diff --git a/zipline/finance/asset_restrictions.py b/zipline/finance/asset_restrictions.py new file mode 100644 index 00000000..c22d5086 --- /dev/null +++ b/zipline/finance/asset_restrictions.py @@ -0,0 +1,220 @@ +import abc +from numpy import vectorize +from functools import partial, reduce +import operator +import pandas as pd +from six import with_metaclass +from collections import namedtuple +from itertools import groupby + +from zipline.utils.enum import enum +from zipline.utils.numpy_utils import vectorized_is_element +from zipline.assets import Asset + + +Restriction = namedtuple( + 'Restriction', ['asset', 'effective_date', 'state'] +) + + +RESTRICTION_STATES = enum( + 'ALLOWED', + 'FROZEN', +) + + +class Restrictions(with_metaclass(abc.ABCMeta)): + """ + Abstract restricted list interface, representing a set of assets that an + algorithm is restricted from trading. + """ + + @abc.abstractmethod + def is_restricted(self, assets, dt): + """ + Is the asset restricted (RestrictionStates.FROZEN) on the given dt? + + Parameters + ---------- + asset : Asset of iterable of Assets + The asset(s) for which we are querying a restriction + dt : pd.Timestamp + The timestamp of the restriction query + + Returns + ------- + is_restricted : bool or pd.Series[bool] indexed by asset + Is the asset or assets restricted on this dt? + + """ + raise NotImplementedError('is_restricted') + + def __or__(self, other_restriction): + """Base implementation for combining two restrictions. + """ + # If the right side is a _UnionRestrictions, defers to the + # _UnionRestrictions implementation of `|`, which intelligently + # flattens restricted lists + if isinstance(other_restriction, _UnionRestrictions): + return other_restriction | self + return _UnionRestrictions([self, other_restriction]) + + +class _UnionRestrictions(Restrictions): + """ + A union of a number of sub restrictions. + + Parameters + ---------- + sub_restrictions : iterable of Restrictions (but not _UnionRestrictions) + The Restrictions to be added together + + Notes + ----- + - Consumers should not construct instances of this class directly, but + instead use the `|` operator to combine restrictions + """ + + def __new__(cls, sub_restrictions): + # Filter out NoRestrictions and deal with resulting cases involving + # one or zero sub_restrictions + sub_restrictions = [ + r for r in sub_restrictions if not isinstance(r, NoRestrictions) + ] + if len(sub_restrictions) == 0: + return NoRestrictions() + elif len(sub_restrictions) == 1: + return sub_restrictions[0] + + new_instance = super(_UnionRestrictions, cls).__new__(cls) + new_instance.sub_restrictions = sub_restrictions + return new_instance + + def __or__(self, other_restriction): + """ + Overrides the base implementation for combining two restrictions, of + which the left side is a _UnionRestrictions. + """ + # Flatten the underlying sub restrictions of _UnionRestrictions + if isinstance(other_restriction, _UnionRestrictions): + new_sub_restrictions = \ + self.sub_restrictions + other_restriction.sub_restrictions + else: + new_sub_restrictions = self.sub_restrictions + [other_restriction] + + return _UnionRestrictions(new_sub_restrictions) + + def is_restricted(self, assets, dt): + if isinstance(assets, Asset): + return any( + r.is_restricted(assets, dt) for r in self.sub_restrictions + ) + + return reduce( + operator.or_, + (r.is_restricted(assets, dt) for r in self.sub_restrictions) + ) + + +class NoRestrictions(Restrictions): + """ + A no-op restrictions that contains no restrictions. + """ + def is_restricted(self, assets, dt): + if isinstance(assets, Asset): + return False + return pd.Series(index=pd.Index(assets), data=False) + + +class StaticRestrictions(Restrictions): + """ + Static restrictions stored in memory that are constant regardless of dt + for each asset. + + Parameters + ---------- + restricted_list : iterable of assets + The assets to be restricted + """ + + def __init__(self, restricted_list): + self._restricted_set = frozenset(restricted_list) + + def is_restricted(self, assets, dt): + """ + An asset is restricted for all dts if it is in the static list. + """ + if isinstance(assets, Asset): + return assets in self._restricted_set + return pd.Series( + index=pd.Index(assets), + data=vectorized_is_element(assets, self._restricted_set) + ) + + +class HistoricalRestrictions(Restrictions): + """ + Historical restrictions stored in memory with effective dates for each + asset. + + Parameters + ---------- + restrictions : iterable of namedtuple Restriction + The restrictions, each defined by an asset, effective date and state + """ + + def __init__(self, restrictions): + # A dict mapping each asset to its restrictions, which are sorted by + # ascending order of effective_date + self._restrictions_by_asset = { + asset: sorted( + restrictions_for_asset, key=lambda x: x.effective_date + ) + for asset, restrictions_for_asset + in groupby(restrictions, lambda x: x.asset) + } + + def is_restricted(self, assets, dt): + """ + Returns whether or not an asset or iterable of assets is restricted + on a dt. + """ + if isinstance(assets, Asset): + return self._is_restricted_for_asset(assets, dt) + + is_restricted = partial(self._is_restricted_for_asset, dt=dt) + return pd.Series( + index=pd.Index(assets), + data=vectorize(is_restricted, otypes=[bool])(assets) + ) + + def _is_restricted_for_asset(self, asset, dt): + state = RESTRICTION_STATES.ALLOWED + for r in self._restrictions_by_asset.get(asset, ()): + if r.effective_date > dt: + break + state = r.state + return state == RESTRICTION_STATES.FROZEN + + +class SecurityListRestrictions(Restrictions): + """ + Restrictions based on a security list. + + Parameters + ---------- + restrictions : zipline.utils.security_list.SecurityList + The restrictions defined by a SecurityList + """ + + def __init__(self, security_list_by_dt): + self.current_securities = security_list_by_dt.current_securities + + def is_restricted(self, assets, dt): + securities_in_list = self.current_securities(dt) + if isinstance(assets, Asset): + return assets in securities_in_list + return pd.Series( + index=pd.Index(assets), + data=vectorized_is_element(assets, securities_in_list) + ) diff --git a/zipline/finance/controls.py b/zipline/finance/controls.py index b1c46de8..ca31a462 100644 --- a/zipline/finance/controls.py +++ b/zipline/finance/controls.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc +import logbook import pandas as pd @@ -23,6 +24,8 @@ from zipline.errors import ( TradingControlViolation, ) +log = logbook.Logger('TradingControl') + class TradingControl(with_metaclass(abc.ABCMeta)): """ @@ -30,11 +33,12 @@ class TradingControl(with_metaclass(abc.ABCMeta)): algorithm. """ - def __init__(self, **kwargs): + def __init__(self, on_error, **kwargs): """ Track any arguments that should be printed in the error message generated by self.fail. """ + self.on_error = on_error self.__fail_args = kwargs @abc.abstractmethod @@ -57,23 +61,36 @@ class TradingControl(with_metaclass(abc.ABCMeta)): """ raise NotImplementedError - def fail(self, asset, amount, datetime, metadata=None): - """ - Raise a TradingControlViolation with information about the failure. - - If dynamic information should be displayed as well, pass it in via - `metadata`. - """ + def _constraint_msg(self, metadata): constraint = repr(self) if metadata: constraint = "{constraint} (Metadata: {metadata})".format( constraint=constraint, metadata=metadata ) - raise TradingControlViolation(asset=asset, - amount=amount, - datetime=datetime, - constraint=constraint) + return constraint + + def handle_violation(self, asset, amount, datetime, metadata=None): + """ + Handle a TradingControlViolation, either by raising or logging and + error with information about the failure. + + If dynamic information should be displayed as well, pass it in via + `metadata`. + """ + constraint = self._constraint_msg(metadata) + + if self.on_error == 'fail': + raise TradingControlViolation( + asset=asset, + amount=amount, + datetime=datetime, + constraint=constraint) + elif self.on_error == 'log': + log.error("Order for {amount} shares of {asset} at {dt} " + "violates trading constraint {constraint}", + amount=amount, asset=asset, dt=datetime, + constraint=constraint) def __repr__(self): return "{name}({attrs})".format(name=self.__class__.__name__, @@ -86,9 +103,9 @@ class MaxOrderCount(TradingControl): placed in a given trading day. """ - def __init__(self, max_count): + def __init__(self, on_error, max_count): - super(MaxOrderCount, self).__init__(max_count=max_count) + super(MaxOrderCount, self).__init__(on_error, max_count=max_count) self.orders_placed = 0 self.max_count = max_count self.current_date = None @@ -96,9 +113,9 @@ class MaxOrderCount(TradingControl): def validate(self, asset, amount, - _portfolio, + portfolio, algo_datetime, - _algo_current_data): + algo_current_data): """ Fail if we've already placed self.max_count orders today. """ @@ -110,7 +127,7 @@ class MaxOrderCount(TradingControl): self.current_date = algo_date if self.orders_placed >= self.max_count: - self.fail(asset, amount, algo_datetime) + self.handle_violation(asset, amount, algo_datetime) self.orders_placed += 1 @@ -120,25 +137,25 @@ class RestrictedListOrder(TradingControl): Parameters ---------- - restricted_list : container[Asset] - The assets that cannot be ordered. + restrictions : zipline.finance.asset_restrictions.Restrictions + Object representing restrictions of a group of assets. """ - def __init__(self, restricted_list): - super(RestrictedListOrder, self).__init__() - self.restricted_list = restricted_list + def __init__(self, on_error, restrictions): + super(RestrictedListOrder, self).__init__(on_error) + self.restrictions = restrictions def validate(self, asset, amount, - _portfolio, - _algo_datetime, - _algo_current_data): + portfolio, + algo_datetime, + algo_current_data): """ Fail if the asset is in the restricted_list. """ - if asset in self.restricted_list: - self.fail(asset, amount, _algo_datetime) + if self.restrictions.is_restricted(asset, algo_datetime): + self.handle_violation(asset, amount, algo_datetime) class MaxOrderSize(TradingControl): @@ -148,8 +165,10 @@ class MaxOrderSize(TradingControl): value. """ - def __init__(self, asset=None, max_shares=None, max_notional=None): - super(MaxOrderSize, self).__init__(asset=asset, + def __init__(self, on_error, asset=None, max_shares=None, + max_notional=None): + super(MaxOrderSize, self).__init__(on_error, + asset=asset, max_shares=max_shares, max_notional=max_notional) self.asset = asset @@ -175,7 +194,7 @@ class MaxOrderSize(TradingControl): asset, amount, portfolio, - _algo_datetime, + algo_datetime, algo_current_data): """ Fail if the magnitude of the given order exceeds either self.max_shares @@ -186,7 +205,7 @@ class MaxOrderSize(TradingControl): return if self.max_shares is not None and abs(amount) > self.max_shares: - self.fail(asset, amount, _algo_datetime) + self.handle_violation(asset, amount, algo_datetime) current_asset_price = algo_current_data.current(asset, "price") order_value = amount * current_asset_price @@ -195,7 +214,7 @@ class MaxOrderSize(TradingControl): abs(order_value) > self.max_notional) if too_much_value: - self.fail(asset, amount, _algo_datetime) + self.handle_violation(asset, amount, algo_datetime) class MaxPositionSize(TradingControl): @@ -204,8 +223,10 @@ class MaxPositionSize(TradingControl): be held by an algo for a given asset. """ - def __init__(self, asset=None, max_shares=None, max_notional=None): - super(MaxPositionSize, self).__init__(asset=asset, + def __init__(self, on_error, asset=None, max_shares=None, + max_notional=None): + super(MaxPositionSize, self).__init__(on_error, + asset=asset, max_shares=max_shares, max_notional=max_notional) self.asset = asset @@ -248,7 +269,7 @@ class MaxPositionSize(TradingControl): too_many_shares = (self.max_shares is not None and abs(shares_post_order) > self.max_shares) if too_many_shares: - self.fail(asset, amount, algo_datetime) + self.handle_violation(asset, amount, algo_datetime) current_price = algo_current_data.current(asset, "price") value_post_order = shares_post_order * current_price @@ -257,7 +278,7 @@ class MaxPositionSize(TradingControl): abs(value_post_order) > self.max_notional) if too_much_value: - self.fail(asset, amount, algo_datetime) + self.handle_violation(asset, amount, algo_datetime) class LongOnly(TradingControl): @@ -265,18 +286,21 @@ class LongOnly(TradingControl): TradingControl representing a prohibition against holding short positions. """ + def __init__(self, on_error): + super(LongOnly, self).__init__(on_error) + def validate(self, asset, amount, portfolio, - _algo_datetime, - _algo_current_data): + algo_datetime, + algo_current_data): """ Fail if we would hold negative shares of asset after completing this order. """ if portfolio.positions[asset].amount + amount < 0: - self.fail(asset, amount, _algo_datetime) + self.handle_violation(asset, amount, algo_datetime) class AssetDateBounds(TradingControl): @@ -285,6 +309,9 @@ class AssetDateBounds(TradingControl): its start_date, or after its end_date. """ + def __init__(self, on_error): + super(AssetDateBounds, self).__init__(on_error) + def validate(self, asset, amount, @@ -308,7 +335,8 @@ class AssetDateBounds(TradingControl): metadata = { 'asset_start_date': normalized_start } - self.fail(asset, amount, algo_datetime, metadata=metadata) + self.handle_violation( + asset, amount, algo_datetime, metadata=metadata) # Fail if the algo has passed this Asset's end_date if asset.end_date: normalized_end = pd.Timestamp(asset.end_date).normalize() @@ -316,7 +344,8 @@ class AssetDateBounds(TradingControl): metadata = { 'asset_end_date': normalized_end } - self.fail(asset, amount, algo_datetime, metadata=metadata) + self.handle_violation( + asset, amount, algo_datetime, metadata=metadata) class AccountControl(with_metaclass(abc.ABCMeta)): diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 2ab816d7..1d5b372f 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -38,7 +38,7 @@ class AlgorithmSimulator(object): } def __init__(self, algo, sim_params, data_portal, clock, benchmark_source, - universe_func): + restrictions, universe_func): # ============== # Simulation @@ -47,6 +47,7 @@ class AlgorithmSimulator(object): self.sim_params = sim_params self.env = algo.trading_environment self.data_portal = data_portal + self.restrictions = restrictions # ============== # Algo Setup @@ -89,6 +90,7 @@ class AlgorithmSimulator(object): simulation_dt_func=self.get_simulation_dt, data_frequency=self.sim_params.data_frequency, trading_calendar=self.algo.trading_calendar, + restrictions=self.restrictions, universe_func=universe_func ) diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index d857245a..5d112f12 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -505,9 +505,22 @@ class SetMaxOrderSizeAlgorithm(TradingAlgorithm): class SetDoNotOrderListAlgorithm(TradingAlgorithm): - def initialize(self, sid=None, restricted_list=None): + def initialize(self, sid=None, restricted_list=None, on_error='fail'): self.order_count = 0 - self.set_do_not_order_list(restricted_list) + self.set_do_not_order_list(restricted_list, on_error) + + +class SetAssetRestrictionsAlgorithm(TradingAlgorithm): + def initialize(self, sid=None, restrictions=None, on_error='fail'): + self.order_count = 0 + self.set_asset_restrictions(restrictions, on_error) + + +class SetMultipleAssetRestrictionsAlgorithm(TradingAlgorithm): + def initialize(self, restrictions1, restrictions2, on_error='fail'): + self.order_count = 0 + self.set_asset_restrictions(restrictions1, on_error) + self.set_asset_restrictions(restrictions2, on_error) class SetMaxOrderCountAlgorithm(TradingAlgorithm): @@ -529,7 +542,7 @@ class SetAssetDateBoundsAlgorithm(TradingAlgorithm): AssetDateBounds() trading control in place. """ def initialize(self): - self.register_trading_control(AssetDateBounds()) + self.register_trading_control(AssetDateBounds(on_error='fail')) def handle_data(algo, data): algo.order(algo.sid(999), 1) diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index f0e2aaa5..3570d255 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -36,8 +36,10 @@ from ..utils.classproperty import classproperty from ..utils.final import FinalMeta, final from .core import tmp_asset_finder, make_simple_equity_info from zipline.assets import Equity, Future +from zipline.finance.asset_restrictions import NoRestrictions from zipline.pipeline import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader +from zipline.protocol import BarData from zipline.utils.calendars import ( get_calendar, register_calendar) @@ -1319,3 +1321,17 @@ class WithResponses(object): self.responses = self.enter_instance_context( responses.RequestsMock(), ) + + +class WithCreateBarData(WithDataPortal): + + CREATE_BARDATA_DATA_FREQUENCY = 'minute' + + def create_bardata(self, simulation_dt_func, restrictions=None): + return BarData( + self.data_portal, + simulation_dt_func, + self.CREATE_BARDATA_DATA_FREQUENCY, + self.trading_calendar, + restrictions or NoRestrictions() + ) diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index a8e35a07..0197ece1 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -25,7 +25,6 @@ from pandas import ( DataFrame, date_range, DatetimeIndex, - DateOffset ) from pandas.tseries.offsets import CustomBusinessDay from zipline.utils.calendars._calendar_helpers import ( @@ -810,36 +809,47 @@ class TradingCalendar(with_metaclass(ABCMeta)): def days_at_time(days, t, tz, day_offset=0): """ - Shift an index of days to time t, interpreted in tz. + Create an index of days at time ``t``, interpreted in timezone ``tz``. - Overwrites any existing tz info on the input. + The returned index is localized to UTC. Parameters ---------- days : DatetimeIndex - The "base" time which we want to change. + An index of dates (represented as midnight). t : datetime.time - The time we want to offset @days by + The time to apply as an offset to each day in ``days``. tz : pytz.timezone - The timezone which these times represent + The timezone to use to interpret ``t``. day_offset : int The number of days we want to offset @days by + + Example + ------- + In the example below, the times switch from 13:45 to 12:45 UTC because + March 13th is the daylight savings transition for US/Eastern. All the + times are still 8:45 when interpreted in US/Eastern. + + >>> import pandas as pd; import datetime; import pprint + >>> dts = pd.date_range('2016-03-12', '2016-03-14') + >>> dts_at_845 = days_at_time(dts, datetime.time(8, 45), 'US/Eastern') + >>> pprint.pprint([str(dt) for dt in dts_at_845]) + ['2016-03-12 13:45:00+00:00', + '2016-03-13 12:45:00+00:00', + '2016-03-14 12:45:00+00:00'] """ if len(days) == 0: return days # Offset days without tz to avoid timezone issues. days = DatetimeIndex(days).tz_localize(None) - days_offset = days + DateOffset(days=day_offset) - - # Shift all days to the target time in the local timezone, then - # convert to UTC. - - # FIXME: Once we're off Pandas 16, see if we can replace DateOffset with - # TimeDelta. - return days_offset.shift( - 1, freq=DateOffset(hour=t.hour, minute=t.minute, second=t.second) - ).tz_localize(tz).tz_convert('UTC') + delta = pd.Timedelta( + days=day_offset, + hours=t.hour, + minutes=t.minute, + seconds=t.second, + ) + return (days + delta).tz_localize(tz).tz_convert('UTC') def holidays_at_time(calendar, start, end, time, tz): diff --git a/zipline/utils/security_list.py b/zipline/utils/security_list.py index 50368b59..2a72952c 100644 --- a/zipline/utils/security_list.py +++ b/zipline/utils/security_list.py @@ -1,3 +1,4 @@ +import warnings from datetime import datetime from os import listdir import os.path @@ -7,6 +8,8 @@ import pytz import zipline from zipline.errors import SymbolNotFound +from zipline.finance.asset_restrictions import SecurityListRestrictions +from zipline.zipline_warnings import ZiplineDeprecationWarning DATE_FORMAT = "%Y%m%d" @@ -38,17 +41,26 @@ class SecurityList(object): return knowledge_dates def __iter__(self): - return iter(self.restricted_list) + warnings.warn( + 'Iterating over security_lists is deprecated. Use ' + '`for sid in .current_securities(dt)` instead.', + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + return iter(self.current_securities(self.current_date())) def __contains__(self, item): - return item in self.restricted_list + warnings.warn( + 'Evaluating inclusion in security_lists is deprecated. Use ' + '`sid in .current_securities(dt)` instead.', + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + return item in self.current_securities(self.current_date()) - @property - def restricted_list(self): - - cd = self.current_date() + def current_securities(self, dt): for kd in self._knowledge_dates: - if cd < kd: + if dt < kd: break if kd in self._cache: self._current_set = self._cache[kd] @@ -103,6 +115,10 @@ class SecurityListSet(object): ) return self._leveraged_etf + @property + def restrict_leveraged_etfs(self): + return SecurityListRestrictions(self.leveraged_etf_list) + def load_from_directory(list_name): """