ENH: _UnionRestrictions for combining multiple Restrictions

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):

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
+        )

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'):

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):
     """

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