From b70084c6bfb92a69d3c596e0632ce029ed833a50 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Wed, 14 Sep 2016 14:03:57 -0400 Subject: [PATCH] ENH: `can_trade` should take restricted list into account Additionally, create an option for a violation of a 'do not order' trading control to log an error instead of failing --- tests/test_algorithm.py | 69 ++++-- tests/test_bar_data.py | 398 +++++++++++++++++++------------- zipline/_protocol.pyx | 11 +- zipline/algorithm.py | 40 +++- zipline/api.py | 10 + zipline/finance/controls.py | 111 +++++---- zipline/gens/tradesimulation.py | 4 +- zipline/test_algorithms.py | 6 +- 8 files changed, 422 insertions(+), 227 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 22884a25..32d5e8bc 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -76,6 +76,11 @@ 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.restrictions import ( + Restriction, + HistoricalRestrictions, + RESTRICTION_STATES, +) from zipline.testing import ( FakeDataPortal, create_daily_df_for_asset, @@ -2789,33 +2794,71 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): 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 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 the restricted list to be one sid for the entire simulation, + # and fail. + rlm = HistoricalRestrictions([ + Restriction( + self.sid, + self.sim_params.start_session, + RESTRICTION_STATES.FROZEN) + ]) + algo = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=rlm, + sim_params=self.sim_params, + env=self.env, + ) self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) + + # if the restricted list is a static list, then use a shim. + rlm = [self.sid] + algo = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=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 = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=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 = 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.assertTrue(algo.could_trade) def test_set_max_order_size(self): diff --git a/tests/test_bar_data.py b/tests/test_bar_data.py index f65a5456..b4b1e170 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.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/zipline/_protocol.pyx b/zipline/_protocol.pyx index fd0c39b2..b15159fd 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.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..9b6df86c 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -76,11 +76,16 @@ from zipline.finance.execution import ( StopOrder, ) from zipline.finance.performance import PerformanceTracker +from zipline.finance.restrictions import Restrictions from zipline.finance.slippage import ( VolumeShareSlippage, SlippageModel ) from zipline.finance.cancel_policy import NeverCancel, CancelPolicy +from zipline.finance.restrictions import ( + NoopRestrictions, + StaticRestrictions +) from zipline.assets import Asset, Future from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.pipeline import Pipeline @@ -120,6 +125,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 +424,8 @@ class TradingAlgorithm(object): # A dictionary of the actual capital change deltas, keyed by timestamp self.capital_change_deltas = {} + self.restrictions = NoopRestrictions() + def init_engine(self, get_loader): """ Construct and store a PipelineEngine from loader. @@ -564,6 +572,7 @@ class TradingAlgorithm(object): self.data_portal, self._create_clock(), self._create_benchmark_source(), + self.restrictions, universe_func=self._calculate_universe ) @@ -2083,7 +2092,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 +2117,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 +2146,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 +2160,32 @@ 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) + + if isinstance(restricted_list, (list, tuple, set)): + restricted_list = StaticRestrictions(restricted_list) + + control = RestrictedListOrder(on_error, restricted_list) self.register_trading_control(control) + self.restrictions = restricted_list @api_method - def set_long_only(self): + 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..426b5f26 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.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/controls.py b/zipline/finance/controls.py index b1c46de8..95d61c89 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.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..052dd821 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -505,9 +505,9 @@ 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 SetMaxOrderCountAlgorithm(TradingAlgorithm): @@ -529,7 +529,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)