From 25ba4369c3ba715387273d4ca2a52c8d00581fc6 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Fri, 30 Sep 2016 12:26:59 -0400 Subject: [PATCH] ENH: _UnionRestrictions for combining multiple Restrictions --- tests/test_algorithm.py | 25 +++++ tests/test_restrictions.py | 136 ++++++++++++++++++++++++++ zipline/algorithm.py | 2 +- zipline/finance/asset_restrictions.py | 65 +++++++++++- zipline/test_algorithms.py | 7 ++ 5 files changed, 233 insertions(+), 2 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index bfe730e1..dea0ab45 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -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): diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py index 64e3f500..9ec90e66 100644 --- a/tests/test_restrictions.py +++ b/tests/test_restrictions.py @@ -12,6 +12,7 @@ from zipline.finance.asset_restrictions import ( StaticRestrictions, SecurityListRestrictions, NoRestrictions, + _UnionRestrictions, ) from zipline.testing import parameter_space @@ -284,3 +285,138 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): 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/zipline/algorithm.py b/zipline/algorithm.py index df437706..dbfb373a 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -2215,7 +2215,7 @@ class TradingAlgorithm(object): """ control = RestrictedListOrder(on_error, restrictions) self.register_trading_control(control) - self.restrictions = restrictions + self.restrictions |= restrictions @api_method def set_long_only(self, on_error='fail'): diff --git a/zipline/finance/asset_restrictions.py b/zipline/finance/asset_restrictions.py index 463eb916..20b7eb73 100644 --- a/zipline/finance/asset_restrictions.py +++ b/zipline/finance/asset_restrictions.py @@ -1,6 +1,7 @@ import abc from numpy import vectorize -from functools import partial +from functools import partial, reduce +import operator import pandas as pd from six import with_metaclass from collections import namedtuple @@ -47,6 +48,68 @@ class Restrictions(with_metaclass(abc.ABCMeta)): """ raise NotImplementedError('is_restricted') + def __or__(self, other_restriction): + """ + Base implementation for combining two restrictions. If the right side + is a _UnionRestrictions, calls the overriding implementation with + _UnionRestrictions on the left side + """ + if isinstance(other_restriction, _UnionRestrictions): + return _UnionRestrictions.__or__(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 + """ + + def __new__(cls, sub_restrictions): + """ + Returns a _UnionRestrictions defined by a list of sub_restrictions, + while dealing with trivial NoRestrictions cases + """ + 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 if the left side is a + _UnionRestrictions. Extracts the underlying sub_restrictions from the + _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): """ diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 41121195..5d112f12 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -516,6 +516,13 @@ class SetAssetRestrictionsAlgorithm(TradingAlgorithm): 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): def initialize(self, count): self.order_count = 0