diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index b9c874e7..403144d2 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -55,6 +55,7 @@ from zipline.test_algorithms import ( SetMaxPositionSizeAlgorithm, SetMaxOrderCountAlgorithm, SetMaxOrderSizeAlgorithm, + SetDoNotOrderListAlgorithm, api_algo, api_get_environment_algo, api_symbol_algo, @@ -955,6 +956,27 @@ class TestTradingControls(TestCase): algo = SetMaxPositionSizeAlgorithm(max_shares=10, max_notional=61.0) 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]) + + def handle_data(algo, data): + algo.order(self.sid, 100) + algo.order_count += 1 + + self.check_algo_fails(algo, handle_data, 0) + + # set the restricted list to exclude the sid, and succeed + algo = SetDoNotOrderListAlgorithm(sid=self.sid, + restricted_list=[134, 135, 136]) + + def handle_data(algo, data): + algo.order(self.sid, 100) + algo.order_count += 1 + + self.check_algo_succeeds(algo, handle_data) + def test_set_max_order_size(self): # Buy one share. diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 64e9b9a0..9a705d86 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -48,6 +48,7 @@ from zipline.finance.controls import ( MaxOrderCount, MaxOrderSize, MaxPositionSize, + RestrictedListOrder ) from zipline.finance.execution import ( LimitOrder, @@ -1080,6 +1081,14 @@ class TradingAlgorithm(object): control = MaxOrderCount(max_count) self.register_trading_control(control) + @api_method + def set_do_not_order_list(self, restricted_list): + """ + Set a restriction on which sids can be ordered. + """ + control = RestrictedListOrder(restricted_list) + self.register_trading_control(control) + @api_method def set_long_only(self): """ diff --git a/zipline/finance/controls.py b/zipline/finance/controls.py index 4e5f6511..d02f4fcf 100644 --- a/zipline/finance/controls.py +++ b/zipline/finance/controls.py @@ -99,6 +99,30 @@ class MaxOrderCount(TradingControl): self.orders_placed += 1 +class RestrictedListOrder(TradingControl): + """ + TradingControl representing a restricted list of securities that + cannot be ordered by the algorithm. + """ + + def __init__(self, restricted_list): + + super(RestrictedListOrder, self).__init__() + self.restricted_list = set(restricted_list) + + def validate(self, + sid, + amount, + _portfolio, + _algo_datetime, + _algo_current_data): + """ + Fail if the sid is in the restricted_list. + """ + if sid in self.restricted_list: + self.fail(sid, amount) + + class MaxOrderSize(TradingControl): """ TradingControl representing a limit on the magnitude of any single order diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index d25adbc3..132b6bb9 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -430,6 +430,12 @@ class SetMaxOrderSizeAlgorithm(TradingAlgorithm): max_notional=max_notional) +class SetDoNotOrderListAlgorithm(TradingAlgorithm): + def initialize(self, sid=None, restricted_list=None): + self.order_count = 0 + self.set_do_not_order_list(restricted_list) + + class SetMaxOrderCountAlgorithm(TradingAlgorithm): def initialize(self, count): self.order_count = 0