mirror of
https://github.com/saymrwulf/zipline.git
synced 2026-05-16 21:10:11 +00:00
419 lines
17 KiB
Python
419 lines
17 KiB
Python
import pandas as pd
|
|
from pandas.util.testing import assert_series_equal
|
|
from six import iteritems
|
|
from functools import partial
|
|
|
|
from toolz import groupby
|
|
|
|
from zipline.finance.asset_restrictions import (
|
|
RESTRICTION_STATES,
|
|
Restriction,
|
|
HistoricalRestrictions,
|
|
StaticRestrictions,
|
|
SecurityListRestrictions,
|
|
NoRestrictions,
|
|
_UnionRestrictions,
|
|
)
|
|
|
|
from zipline.testing import parameter_space
|
|
from zipline.testing.fixtures import (
|
|
WithDataPortal,
|
|
ZiplineTestCase,
|
|
)
|
|
|
|
|
|
def str_to_ts(dt_str):
|
|
return pd.Timestamp(dt_str, tz='UTC')
|
|
|
|
|
|
FROZEN = RESTRICTION_STATES.FROZEN
|
|
ALLOWED = RESTRICTION_STATES.ALLOWED
|
|
MINUTE = pd.Timedelta(minutes=1)
|
|
|
|
|
|
class RestrictionsTestCase(WithDataPortal, ZiplineTestCase):
|
|
|
|
ASSET_FINDER_EQUITY_SIDS = 1, 2, 3
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(RestrictionsTestCase, cls).init_class_fixtures()
|
|
cls.ASSET1 = cls.asset_finder.retrieve_asset(1)
|
|
cls.ASSET2 = cls.asset_finder.retrieve_asset(2)
|
|
cls.ASSET3 = cls.asset_finder.retrieve_asset(3)
|
|
cls.ALL_ASSETS = [cls.ASSET1, cls.ASSET2, cls.ASSET3]
|
|
|
|
def assert_is_restricted(self, rl, asset, dt):
|
|
self.assertTrue(rl.is_restricted(asset, dt))
|
|
|
|
def assert_not_restricted(self, rl, asset, dt):
|
|
self.assertFalse(rl.is_restricted(asset, dt))
|
|
|
|
def assert_all_restrictions(self, rl, expected, dt):
|
|
self.assert_many_restrictions(rl, self.ALL_ASSETS, expected, dt)
|
|
|
|
def assert_many_restrictions(self, rl, assets, expected, dt):
|
|
assert_series_equal(
|
|
rl.is_restricted(assets, dt),
|
|
pd.Series(index=pd.Index(assets), data=expected),
|
|
)
|
|
|
|
@parameter_space(
|
|
date_offset=(
|
|
pd.Timedelta(0),
|
|
pd.Timedelta('1 minute'),
|
|
pd.Timedelta('15 hours 5 minutes')
|
|
),
|
|
restriction_order=(
|
|
list(range(6)), # Keep restrictions in order.
|
|
[0, 2, 1, 3, 5, 4], # Re-order within asset.
|
|
[0, 3, 1, 4, 2, 5], # Scramble assets, maintain per-asset order.
|
|
[0, 5, 2, 3, 1, 4], # Scramble assets and per-asset order.
|
|
),
|
|
__fail_fast=True,
|
|
)
|
|
def test_historical_restrictions(self, date_offset, restriction_order):
|
|
"""
|
|
Test historical restrictions for both interday and intraday
|
|
restrictions, as well as restrictions defined in/not in order, for both
|
|
single- and multi-asset queries
|
|
"""
|
|
def rdate(s):
|
|
"""Convert a date string into a restriction for that date."""
|
|
# Add date_offset to check that we handle intraday changes.
|
|
return str_to_ts(s) + date_offset
|
|
|
|
base_restrictions = [
|
|
Restriction(self.ASSET1, rdate('2011-01-04'), FROZEN),
|
|
Restriction(self.ASSET1, rdate('2011-01-05'), ALLOWED),
|
|
Restriction(self.ASSET1, rdate('2011-01-06'), FROZEN),
|
|
Restriction(self.ASSET2, rdate('2011-01-05'), FROZEN),
|
|
Restriction(self.ASSET2, rdate('2011-01-06'), ALLOWED),
|
|
Restriction(self.ASSET2, rdate('2011-01-07'), FROZEN),
|
|
]
|
|
# Scramble the restrictions based on restriction_order to check that we
|
|
# don't depend on the order in which restrictions are provided to us.
|
|
all_restrictions = [base_restrictions[i] for i in restriction_order]
|
|
|
|
restrictions_by_asset = groupby(lambda r: r.asset, all_restrictions)
|
|
|
|
rl = HistoricalRestrictions(all_restrictions)
|
|
assert_not_restricted = partial(self.assert_not_restricted, rl)
|
|
assert_is_restricted = partial(self.assert_is_restricted, rl)
|
|
assert_all_restrictions = partial(self.assert_all_restrictions, rl)
|
|
|
|
# Check individual restrictions.
|
|
for asset, r_history in iteritems(restrictions_by_asset):
|
|
freeze_dt, unfreeze_dt, re_freeze_dt = (
|
|
sorted([r.effective_date for r in r_history])
|
|
)
|
|
|
|
# Starts implicitly unrestricted. Restricted on or after the freeze
|
|
assert_not_restricted(asset, freeze_dt - MINUTE)
|
|
assert_is_restricted(asset, freeze_dt)
|
|
assert_is_restricted(asset, freeze_dt + MINUTE)
|
|
|
|
# Unrestricted on or after the unfreeze
|
|
assert_is_restricted(asset, unfreeze_dt - MINUTE)
|
|
assert_not_restricted(asset, unfreeze_dt)
|
|
assert_not_restricted(asset, unfreeze_dt + MINUTE)
|
|
|
|
# Restricted again on or after the freeze
|
|
assert_not_restricted(asset, re_freeze_dt - MINUTE)
|
|
assert_is_restricted(asset, re_freeze_dt)
|
|
assert_is_restricted(asset, re_freeze_dt + MINUTE)
|
|
|
|
# Should stay restricted for the rest of time
|
|
assert_is_restricted(asset, re_freeze_dt + MINUTE * 1000000)
|
|
|
|
# Check vectorized restrictions.
|
|
# Expected results for [self.ASSET1, self.ASSET2, self.ASSET3],
|
|
# ASSET3 is always False as it has no defined restrictions
|
|
|
|
# 01/04 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: ALLOWED
|
|
d0 = rdate('2011-01-04')
|
|
assert_all_restrictions([False, False, False], d0 - MINUTE)
|
|
assert_all_restrictions([True, False, False], d0)
|
|
assert_all_restrictions([True, False, False], d0 + MINUTE)
|
|
|
|
# 01/05 XX:00 ASSET1: FROZEN --> ALLOWED; ASSET2: ALLOWED --> FROZEN
|
|
d1 = rdate('2011-01-05')
|
|
assert_all_restrictions([True, False, False], d1 - MINUTE)
|
|
assert_all_restrictions([False, True, False], d1)
|
|
assert_all_restrictions([False, True, False], d1 + MINUTE)
|
|
|
|
# 01/06 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: FROZEN --> ALLOWED
|
|
d2 = rdate('2011-01-06')
|
|
assert_all_restrictions([False, True, False], d2 - MINUTE)
|
|
assert_all_restrictions([True, False, False], d2)
|
|
assert_all_restrictions([True, False, False], d2 + MINUTE)
|
|
|
|
# 01/07 XX:00 ASSET1: FROZEN; ASSET2: ALLOWED --> FROZEN
|
|
d3 = rdate('2011-01-07')
|
|
assert_all_restrictions([True, False, False], d3 - MINUTE)
|
|
assert_all_restrictions([True, True, False], d3)
|
|
assert_all_restrictions([True, True, False], d3 + MINUTE)
|
|
|
|
# Should stay restricted for the rest of time
|
|
assert_all_restrictions(
|
|
[True, True, False],
|
|
d3 + (MINUTE * 10000000)
|
|
)
|
|
|
|
def test_historical_restrictions_consecutive_states(self):
|
|
"""
|
|
Test that defining redundant consecutive restrictions still works
|
|
"""
|
|
rl = HistoricalRestrictions([
|
|
Restriction(self.ASSET1, str_to_ts('2011-01-04'), ALLOWED),
|
|
Restriction(self.ASSET1, str_to_ts('2011-01-05'), ALLOWED),
|
|
Restriction(self.ASSET1, str_to_ts('2011-01-06'), FROZEN),
|
|
Restriction(self.ASSET1, str_to_ts('2011-01-07'), FROZEN),
|
|
])
|
|
|
|
assert_not_restricted = partial(self.assert_not_restricted, rl)
|
|
assert_is_restricted = partial(self.assert_is_restricted, rl)
|
|
|
|
# (implicit) ALLOWED --> ALLOWED
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04') - MINUTE)
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04'))
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04') + MINUTE)
|
|
|
|
# ALLOWED --> ALLOWED
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05') - MINUTE)
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05'))
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05') + MINUTE)
|
|
|
|
# ALLOWED --> FROZEN
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-06') - MINUTE)
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-06'))
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-06') + MINUTE)
|
|
|
|
# FROZEN --> FROZEN
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07') - MINUTE)
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07'))
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07') + MINUTE)
|
|
|
|
def test_static_restrictions(self):
|
|
"""
|
|
Test single- and multi-asset queries on static restrictions
|
|
"""
|
|
|
|
restricted_a1 = self.ASSET1
|
|
restricted_a2 = self.ASSET2
|
|
unrestricted_a3 = self.ASSET3
|
|
|
|
rl = StaticRestrictions([restricted_a1, restricted_a2])
|
|
assert_not_restricted = partial(self.assert_not_restricted, rl)
|
|
assert_is_restricted = partial(self.assert_is_restricted, rl)
|
|
assert_all_restrictions = partial(self.assert_all_restrictions, rl)
|
|
|
|
for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03',
|
|
'2011-01-04',
|
|
'2011-01-04 1:01',
|
|
'2020-01-04')]:
|
|
assert_is_restricted(restricted_a1, dt)
|
|
assert_is_restricted(restricted_a2, dt)
|
|
assert_not_restricted(unrestricted_a3, dt)
|
|
|
|
assert_all_restrictions([True, True, False], dt)
|
|
|
|
def test_security_list_restrictions(self):
|
|
"""
|
|
Test single- and multi-asset queries on restrictions defined by
|
|
zipline.utils.security_list.SecurityList
|
|
"""
|
|
|
|
# A mock SecurityList object filled with fake data
|
|
class SecurityList(object):
|
|
def __init__(self, assets_by_dt):
|
|
self.assets_by_dt = assets_by_dt
|
|
|
|
def current_securities(self, dt):
|
|
return self.assets_by_dt[dt]
|
|
|
|
assets_by_dt = {
|
|
str_to_ts('2011-01-03'): [self.ASSET1],
|
|
str_to_ts('2011-01-04'): [self.ASSET2, self.ASSET3],
|
|
str_to_ts('2011-01-05'): [self.ASSET1, self.ASSET2, self.ASSET3],
|
|
}
|
|
|
|
rl = SecurityListRestrictions(SecurityList(assets_by_dt))
|
|
|
|
assert_not_restricted = partial(self.assert_not_restricted, rl)
|
|
assert_is_restricted = partial(self.assert_is_restricted, rl)
|
|
assert_all_restrictions = partial(self.assert_all_restrictions, rl)
|
|
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-03'))
|
|
assert_not_restricted(self.ASSET2, str_to_ts('2011-01-03'))
|
|
assert_not_restricted(self.ASSET3, str_to_ts('2011-01-03'))
|
|
assert_all_restrictions(
|
|
[True, False, False], str_to_ts('2011-01-03')
|
|
)
|
|
|
|
assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04'))
|
|
assert_is_restricted(self.ASSET2, str_to_ts('2011-01-04'))
|
|
assert_is_restricted(self.ASSET3, str_to_ts('2011-01-04'))
|
|
assert_all_restrictions(
|
|
[False, True, True], str_to_ts('2011-01-04')
|
|
)
|
|
|
|
assert_is_restricted(self.ASSET1, str_to_ts('2011-01-05'))
|
|
assert_is_restricted(self.ASSET2, str_to_ts('2011-01-05'))
|
|
assert_is_restricted(self.ASSET3, str_to_ts('2011-01-05'))
|
|
assert_all_restrictions(
|
|
[True, True, True],
|
|
str_to_ts('2011-01-05')
|
|
)
|
|
|
|
def test_noop_restrictions(self):
|
|
"""
|
|
Test single- and multi-asset queries on no-op restrictions
|
|
"""
|
|
|
|
rl = NoRestrictions()
|
|
assert_not_restricted = partial(self.assert_not_restricted, rl)
|
|
assert_all_restrictions = partial(self.assert_all_restrictions, rl)
|
|
|
|
for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03',
|
|
'2011-01-04',
|
|
'2020-01-04')]:
|
|
assert_not_restricted(self.ASSET1, dt)
|
|
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
|
|
)
|