ENH: _UnionRestrictions for combining multiple Restrictions

This commit is contained in:
Andrew Liang
2016-09-30 12:26:59 -04:00
parent 148d2a5273
commit 25ba4369c3
5 changed files with 233 additions and 2 deletions
+25
View File
@@ -129,6 +129,7 @@ from zipline.test_algorithms import (
SetMaxOrderSizeAlgorithm,
SetDoNotOrderListAlgorithm,
SetAssetRestrictionsAlgorithm,
SetMultipleAssetRestrictionsAlgorithm,
SetMaxLeverageAlgorithm,
api_algo,
api_get_environment_algo,
@@ -2862,6 +2863,30 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase):
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):