From 0119aba410174f3f1a3d6d79427a1d121649b6d1 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Wed, 14 Sep 2016 14:01:23 -0400 Subject: [PATCH 01/15] ENH: A point-in-time restricted list with restrictions stored in memory An ABC Restrictions defines a group of restrictions responsible for returning restriction information for sids on certain dts. An InMemoryRestrictions is a point-in-time group of such restrictions, with all restrictions and their dates passed in upon instantiation. A StaticRestrictedList takes a list of sids, restricting them at all dates --- tests/test_restrictions.py | 230 ++++++++++++++++++++++++++++++++ zipline/finance/restrictions.py | 129 ++++++++++++++++++ 2 files changed, 359 insertions(+) create mode 100644 tests/test_restrictions.py create mode 100644 zipline/finance/restrictions.py diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py new file mode 100644 index 00000000..c1817b68 --- /dev/null +++ b/tests/test_restrictions.py @@ -0,0 +1,230 @@ +import pandas as pd +from pandas.util.testing import assert_series_equal +from nose_parameterized import parameterized +from six import iteritems +from functools import partial + +from zipline.finance.restrictions import ( + RESTRICTION_STATES, + Restriction, + HistoricalRestrictions, + StaticRestrictions, + NoopRestrictions, +) + +from zipline.testing.fixtures import ( + WithDataPortal, + ZiplineTestCase, +) + +str_to_ts = lambda dt_str: 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) + + 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_vectorized_results(self, rl, expected, dt): + assert_series_equal( + rl.is_restricted([self.ASSET1, self.ASSET2, self.ASSET3], dt), + pd.Series( + index=pd.Index([self.ASSET1, self.ASSET2, self.ASSET3]), + data=expected + ) + ) + + @parameterized.expand([ + ('_'.join([timing, ordering]), + timing == 'intraday', + ordering == 'ordered') + for timing in ['intraday', 'interday'] + for ordering in ['ordered', 'unordered'] + ]) + def test_historical_restrictions(self, name, is_intraday, is_ordered): + """ + 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 + """ + + hour_of_day = ' 15:00' if is_intraday else '' + + if is_ordered: + restriction_dates = { + self.ASSET1: [ + (str_to_ts('2011-01-04' + hour_of_day), FROZEN), + (str_to_ts('2011-01-05' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-06' + hour_of_day), FROZEN), + ], + self.ASSET2: [ + (str_to_ts('2011-01-05' + hour_of_day), FROZEN), + (str_to_ts('2011-01-06' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-07' + hour_of_day), FROZEN), + ], + } + else: + restriction_dates = { + self.ASSET1: [ + (str_to_ts('2011-01-05' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-06' + hour_of_day), FROZEN), + (str_to_ts('2011-01-04' + hour_of_day), FROZEN), + ], + self.ASSET2: [ + (str_to_ts('2011-01-06' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-05' + hour_of_day), FROZEN), + (str_to_ts('2011-01-07' + hour_of_day), FROZEN), + ], + } + + restrictions = sum([ + [Restriction(asset, info[0], info[1]) for info in r_history] + for asset, r_history in iteritems(restriction_dates) + ], []) + rl = HistoricalRestrictions(restrictions) + + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + assert_vectorized_results = partial(self.assert_vectorized_results, rl) + + for asset, r_history in iteritems(restriction_dates): + dts = sorted([info[0] for info in r_history]) + + # Not restricted until on or after the freeze + assert_not_restricted(asset, dts[0] - MINUTE) + assert_is_restricted(asset, dts[0]) + assert_is_restricted(asset, dts[0] + MINUTE) + + # Unrestricted on or after the unfreeze + assert_is_restricted(asset, dts[1] - MINUTE) + assert_not_restricted(asset, dts[1]) + assert_not_restricted(asset, dts[1] + MINUTE) + + # Restricted again on or after the freeze + assert_not_restricted(asset, dts[2] - MINUTE) + assert_is_restricted(asset, dts[2]) + assert_is_restricted(asset, dts[2] + MINUTE) + # Should stay restricted for the rest of time + assert_is_restricted(asset, dts[2] + MINUTE * 1000000) + + dts = [str_to_ts(ts + hour_of_day) for ts in ['2011-01-04', + '2011-01-05', + '2011-01-06', + '2011-01-07']] + + # 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 + assert_vectorized_results([False, False, False], dts[0] - MINUTE) + assert_vectorized_results([True, False, False], dts[0]) + assert_vectorized_results([True, False, False], dts[0] + MINUTE) + + # 01/05 XX:00 ASSET1: FROZEN --> ALLOWED; ASSET2: ALLOWED --> FROZEN + assert_vectorized_results([True, False, False], dts[1] - MINUTE) + assert_vectorized_results([False, True, False], dts[1]) + assert_vectorized_results([False, True, False], dts[1] + MINUTE) + + # 01/06 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: FROZEN --> ALLOWED + assert_vectorized_results([False, True, False], dts[2] - MINUTE) + assert_vectorized_results([True, False, False], dts[2]) + assert_vectorized_results([True, False, False], dts[2] + MINUTE) + + # 01/07 XX:00 ASSET1: FROZEN; ASSET2: ALLOWED --> FROZEN + assert_vectorized_results([True, False, False], dts[3] - MINUTE) + assert_vectorized_results([True, True, False], dts[3]) + assert_vectorized_results([True, True, False], dts[3] + MINUTE) + # Should stay restricted for the rest of time + assert_vectorized_results( + [True, True, False], + dts[3] + 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_vectorized_results = partial(self.assert_vectorized_results, rl) + + for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03', + '2011-01-04', + '2020-01-04')]: + assert_is_restricted(restricted_a1, dt) + assert_is_restricted(restricted_a2, dt) + assert_not_restricted(unrestricted_a3, dt) + + assert_vectorized_results([True, True, False], dt) + + def test_noop_restrictions(self): + """ + Test single- and multi-asset queries on no-op restrictions + """ + + rl = NoopRestrictions() + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_vectorized_results = partial(self.assert_vectorized_results, 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_vectorized_results([False, False, False], dt) diff --git a/zipline/finance/restrictions.py b/zipline/finance/restrictions.py new file mode 100644 index 00000000..e977d2d8 --- /dev/null +++ b/zipline/finance/restrictions.py @@ -0,0 +1,129 @@ +import abc +from numpy import vectorize +from functools import partial +import pandas as pd +from six import with_metaclass +from collections import namedtuple +from itertools import groupby + +from zipline.utils.enum import enum +from zipline.utils.numpy_utils import vectorized_is_element +from zipline.assets import Asset + + +Restriction = namedtuple( + 'Restriction', ['asset', 'effective_date', 'state'] +) + + +RESTRICTION_STATES = enum( + 'ALLOWED', + 'FROZEN', +) + + +class Restrictions(with_metaclass(abc.ABCMeta)): + """ + Abstract restricted list interface + """ + + @abc.abstractmethod + def is_restricted(self, assets, dt): + """ + Is the asset restricted (RestrictionStates.FROZEN) on the given dt? + + Parameters + ---------- + asset : Asset of iterable of Assets + The asset(s) for which we are querying a restriction + dt : pd.Timestamp + The timestamp of the restriction query + + Returns + ------- + is_restricted : bool or pd.Series[bool] indexed by asset + Is the asset or assets restricted on this dt? + + """ + raise NotImplementedError('is_restricted') + + +class NoopRestrictions(Restrictions): + """ + A no-op restrictions that contains no restrictions + """ + def is_restricted(self, assets, dt): + if isinstance(assets, Asset): + return False + return pd.Series(index=pd.Index(assets), data=[False]*len(assets)) + + +class StaticRestrictions(Restrictions): + """ + Static restrictions stored in memory that are constant regardless of dt + for each asset + + Parameters + ---------- + restricted_list : iterable of assets + The assets to be restricted + """ + + def __init__(self, restricted_list): + self._restricted_set = frozenset(restricted_list) + + def is_restricted(self, assets, dt): + """ + An asset is restricted for all dts if it is in the static list + """ + if isinstance(assets, Asset): + return assets in self._restricted_set + return pd.Series( + index=pd.Index(assets), + data=vectorized_is_element(assets, self._restricted_set) + ) + + +class HistoricalRestrictions(Restrictions): + """ + Historical restrictions stored in memory with effective dates for each + asset + + Parameters + ---------- + restrictions : iterable of namedtuple Restriction + The restrictions, each defined by an asset, effective date and state + """ + + def __init__(self, restrictions): + # A dict mapping each asset to its restrictions, which are sorted by + # ascending order of effective_date + self._restrictions_by_asset = { + asset: sorted( + restrictions_for_asset, key=lambda x: x.effective_date + ) + for asset, restrictions_for_asset + in groupby(restrictions, lambda x: x.asset) + } + + def is_restricted(self, assets, dt): + """ + Returns whether or not an asset or iterable of assets is restricted + on a dt + """ + if isinstance(assets, Asset): + return self._is_restricted_for_asset(assets, dt) + + is_restricted = partial(self._is_restricted_for_asset, dt=dt) + return pd.Series( + index=pd.Index(assets), + data=vectorize(is_restricted, otypes=[bool])(assets) + ) + + def _is_restricted_for_asset(self, asset, dt): + state = RESTRICTION_STATES.ALLOWED + for r in self._restrictions_by_asset.get(asset, ()): + if r.effective_date > dt: + break + state = r.state + return state == RESTRICTION_STATES.FROZEN From b70084c6bfb92a69d3c596e0632ce029ed833a50 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Wed, 14 Sep 2016 14:03:57 -0400 Subject: [PATCH 02/15] ENH: `can_trade` should take restricted list into account Additionally, create an option for a violation of a 'do not order' trading control to log an error instead of failing --- tests/test_algorithm.py | 69 ++++-- tests/test_bar_data.py | 398 +++++++++++++++++++------------- zipline/_protocol.pyx | 11 +- zipline/algorithm.py | 40 +++- zipline/api.py | 10 + zipline/finance/controls.py | 111 +++++---- zipline/gens/tradesimulation.py | 4 +- zipline/test_algorithms.py | 6 +- 8 files changed, 422 insertions(+), 227 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 22884a25..32d5e8bc 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -76,6 +76,11 @@ from zipline.finance.commission import PerShare from zipline.finance.execution import LimitOrder from zipline.finance.order import ORDER_STATUS from zipline.finance.trading import SimulationParameters +from zipline.finance.restrictions import ( + Restriction, + HistoricalRestrictions, + RESTRICTION_STATES, +) from zipline.testing import ( FakeDataPortal, create_daily_df_for_asset, @@ -2789,33 +2794,71 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): 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], - sim_params=self.sim_params, - env=self.env, - ) def handle_data(algo, data): + algo.could_trade = data.can_trade(algo.sid(self.sid)) algo.order(algo.sid(self.sid), 100) algo.order_count += 1 + # set the restricted list to be one sid for the entire simulation, + # and fail. + rlm = HistoricalRestrictions([ + Restriction( + self.sid, + self.sim_params.start_session, + RESTRICTION_STATES.FROZEN) + ]) + algo = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=rlm, + sim_params=self.sim_params, + env=self.env, + ) self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) + + # if the restricted list is a static list, then use a shim. + rlm = [self.sid] + algo = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=rlm, + sim_params=self.sim_params, + env=self.env, + ) + self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) + + # just log an error on the violation if we choose not to fail. + algo = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=rlm, + sim_params=self.sim_params, + env=self.env, + on_error='log' + ) + with make_test_handler(self) as log_catcher: + self.check_algo_succeeds(algo, handle_data) + logs = [r.message for r in log_catcher.records] + self.assertIn("Order for 100 shares of Equity(133 [A]) at " + "2006-01-03 21:00:00+00:00 violates trading constraint " + "RestrictedListOrder({})", logs) + self.assertFalse(algo.could_trade) # set the restricted list to exclude the sid, and succeed + rlm = HistoricalRestrictions([ + Restriction( + sid, + self.sim_params.start_session, + RESTRICTION_STATES.FROZEN) for sid in [134, 135, 136] + ]) algo = SetDoNotOrderListAlgorithm( sid=self.sid, - restricted_list=[134, 135, 136], + restricted_list=rlm, sim_params=self.sim_params, env=self.env, ) - - def handle_data(algo, data): - algo.order(algo.sid(self.sid), 100) - algo.order_count += 1 - self.check_algo_succeeds(algo, handle_data) + self.assertTrue(algo.could_trade) def test_set_max_order_size(self): diff --git a/tests/test_bar_data.py b/tests/test_bar_data.py index f65a5456..b4b1e170 100644 --- a/tests/test_bar_data.py +++ b/tests/test_bar_data.py @@ -23,8 +23,11 @@ import pandas as pd from zipline._protocol import handle_non_market_minutes -from zipline.data.data_portal import DataPortal -from zipline.protocol import BarData +from zipline.finance.restrictions import ( + Restriction, + HistoricalRestrictions, + RESTRICTION_STATES, +) from zipline.testing import ( MockDailyBarReader, create_daily_df_for_asset, @@ -32,6 +35,7 @@ from zipline.testing import ( str_to_seconds, ) from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, ZiplineTestCase, ) @@ -49,6 +53,8 @@ field_info = { "close": 0 } +str_to_ts = lambda dt_str: pd.Timestamp(dt_str, tz='UTC') + class WithBarDataChecks(object): def assert_same(self, val1, val2): @@ -95,7 +101,8 @@ class WithBarDataChecks(object): getattr(bar_data, field) -class TestMinuteBarData(WithBarDataChecks, +class TestMinuteBarData(WithCreateBarData, + WithBarDataChecks, WithDataPortal, ZiplineTestCase): START_DATE = pd.Timestamp('2016-01-05', tz='UTC') @@ -205,8 +212,9 @@ class TestMinuteBarData(WithBarDataChecks, # this entire day is before either asset has started trading for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.check_internal_consistency(bar_data) self.assertFalse(bar_data.can_trade(self.ASSET1)) @@ -248,8 +256,9 @@ class TestMinuteBarData(WithBarDataChecks, # this test covers the "IPO morning" case, because asset2 only # has data starting on the 10th minute. - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.check_internal_consistency(bar_data) asset2_has_data = (((idx + 1) % 10) == 0) @@ -328,8 +337,9 @@ class TestMinuteBarData(WithBarDataChecks, # this is the last day the assets exist for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertTrue(bar_data.can_trade(self.ASSET1)) self.assertTrue(bar_data.can_trade(self.ASSET2)) @@ -347,8 +357,9 @@ class TestMinuteBarData(WithBarDataChecks, # this entire day is after both assets have stopped trading for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertFalse(bar_data.can_trade(self.ASSET1)) self.assertFalse(bar_data.can_trade(self.ASSET2)) @@ -390,8 +401,9 @@ class TestMinuteBarData(WithBarDataChecks, ) for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertEqual( idx + 1, bar_data.current(self.SPLIT_ASSET, "price") @@ -408,16 +420,16 @@ class TestMinuteBarData(WithBarDataChecks, ) for idx, minute in enumerate(day0_minutes[-10:-1]): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) self.assertEqual( 380, bar_data.current(self.ILLIQUID_SPLIT_ASSET, "price") ) - bar_data = BarData( - self.data_portal, lambda: day0_minutes[-1], "minute", - self.trading_calendar + bar_data = self.create_bardata( + lambda: day0_minutes[-1], ) self.assertEqual( @@ -426,8 +438,9 @@ class TestMinuteBarData(WithBarDataChecks, ) for idx, minute in enumerate(day1_minutes[0:9]): - bar_data = BarData(self.data_portal, lambda: minute, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) # should be half of 390, due to the split self.assertEqual( @@ -446,12 +459,12 @@ class TestMinuteBarData(WithBarDataChecks, tz='US/Eastern' ) - bar_data = BarData(self.data_portal, lambda: day, "minute", - self.trading_calendar) - bar_data2 = BarData(self.data_portal, - lambda: eight_fortyfive_am_eastern, - "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: day, + ) + bar_data2 = self.create_bardata( + lambda: eight_fortyfive_am_eastern, + ) with handle_non_market_minutes(bar_data), \ handle_non_market_minutes(bar_data2): @@ -482,20 +495,10 @@ class TestMinuteBarData(WithBarDataChecks, def test_get_value_during_non_market_hours(self): # make sure that if we try to get the OHLCV values of ASSET1 during # non-market hours, we don't get the previous market minute's values - futures_cal = get_calendar("us_futures") - data_portal = DataPortal( - self.env.asset_finder, - futures_cal, - first_trading_day=self.DATA_PORTAL_FIRST_TRADING_DAY, - equity_minute_reader=self.bcolz_equity_minute_bar_reader, - ) - - bar_data = BarData( - data_portal, - lambda: pd.Timestamp("2016-01-06 3:15", tz="US/Eastern"), - "minute", - futures_cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: + pd.Timestamp("2016-01-06 4:15", tz="US/Eastern"), ) self.assertTrue(np.isnan(bar_data.current(self.ASSET1, "open"))) @@ -508,14 +511,14 @@ class TestMinuteBarData(WithBarDataChecks, self.assertEqual(390, bar_data.current(self.ASSET1, "price")) def test_can_trade_equity_same_cal_outside_lifetime(self): - cal = get_calendar(self.ASSET1.exchange) # verify that can_trade returns False for the session before the # asset's first session - session_before_asset1_start = cal.previous_session_label( - self.ASSET1.start_date - ) - minutes_for_session = cal.minutes_for_session( + session_before_asset1_start = \ + self.trading_calendar.previous_session_label( + self.ASSET1.start_date + ) + minutes_for_session = self.trading_calendar.minutes_for_session( session_before_asset1_start ) @@ -526,14 +529,14 @@ class TestMinuteBarData(WithBarDataChecks, ) for minute in minutes_to_check: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertFalse(bar_data.can_trade(self.ASSET1)) # after asset lifetime - session_after_asset1_end = cal.next_session_label( + session_after_asset1_end = self.trading_calendar.next_session_label( self.ASSET1.end_date ) bts_after_asset1_end = session_after_asset1_end.replace( @@ -541,32 +544,32 @@ class TestMinuteBarData(WithBarDataChecks, ).tz_convert(None).tz_localize("US/Eastern") minutes_to_check = chain( - cal.minutes_for_session(session_after_asset1_end), + self.trading_calendar.minutes_for_session( + session_after_asset1_end + ), [bts_after_asset1_end] ) for minute in minutes_to_check: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertFalse(bar_data.can_trade(self.ASSET1)) def test_can_trade_equity_same_cal_exchange_closed(self): - cal = get_calendar(self.ASSET1.exchange) - # verify that can_trade returns true for minutes that are # outside the asset's calendar (assuming the asset is alive and # there is a last price), because the asset is alive on the # next market minute. - minutes = cal.minutes_for_sessions_in_range( + minutes = self.trading_calendar.minutes_for_sessions_in_range( self.ASSET1.start_date, self.ASSET1.end_date ) for minute in minutes: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertTrue(bar_data.can_trade(self.ASSET1)) @@ -576,13 +579,13 @@ class TestMinuteBarData(WithBarDataChecks, # 2016-01-05 15:20:00+00:00. Make sure that can_trade returns false # for all minutes in that session before the first trade, and true # for all minutes afterwards. - cal = get_calendar(self.ASSET1.exchange) - minutes_in_session = cal.minutes_for_session(self.ASSET1.start_date) + minutes_in_session = \ + self.trading_calendar.minutes_for_session(self.ASSET1.start_date) for minute in minutes_in_session[0:49]: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertFalse(bar_data.can_trade( @@ -590,14 +593,139 @@ class TestMinuteBarData(WithBarDataChecks, ) for minute in minutes_in_session[50:]: - bar_data = BarData( - self.data_portal, lambda: minute, "minute", cal + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, ) self.assertTrue(bar_data.can_trade( self.HILARIOUSLY_ILLIQUID_ASSET) ) + def test_is_stale_during_non_market_hours(self): + bar_data = self.create_bardata( + lambda: self.equity_minute_bar_days[1], + ) + + with handle_non_market_minutes(bar_data): + self.assertTrue(bar_data.is_stale(self.HILARIOUSLY_ILLIQUID_ASSET)) + + def test_overnight_adjustments(self): + # verify there is a split for SPLIT_ASSET + splits = self.adjustment_reader.get_adjustments_for_sid( + "splits", + self.SPLIT_ASSET.sid + ) + + self.assertEqual(1, len(splits)) + split = splits[0] + self.assertEqual( + split[0], + pd.Timestamp("2016-01-06", tz='UTC') + ) + + # Current day is 1/06/16 + day = self.equity_daily_bar_days[1] + eight_fortyfive_am_eastern = \ + pd.Timestamp("{0}-{1}-{2} 8:45".format( + day.year, day.month, day.day), + tz='US/Eastern' + ) + + bar_data = self.create_bardata( + lambda: eight_fortyfive_am_eastern, + ) + + expected = { + 'open': 391 / 2.0, + 'high': 392 / 2.0, + 'low': 389 / 2.0, + 'close': 390 / 2.0, + 'volume': 39000 * 2.0, + 'price': 390 / 2.0, + } + + with handle_non_market_minutes(bar_data): + for field in OHLCP + ['volume']: + value = bar_data.current(self.SPLIT_ASSET, field) + + # Assert the price is adjusted for the overnight split + self.assertEqual(value, expected[field]) + + def test_can_trade_restricted(self): + """ + Test that can_trade will return False for a sid if it is restricted + on that dt + """ + + minutes_to_check = [ + (str_to_ts("2016-01-05 14:31"), False), + (str_to_ts("2016-01-06 14:31"), False), + (str_to_ts("2016-01-07 14:31"), True), + (str_to_ts("2016-01-07 15:00"), False), + (str_to_ts("2016-01-07 15:30"), True), + ] + + rlm = HistoricalRestrictions([ + Restriction(1, str_to_ts('2016-01-05'), + RESTRICTION_STATES.FROZEN), + Restriction(1, str_to_ts('2016-01-07'), + RESTRICTION_STATES.ALLOWED), + Restriction(1, str_to_ts('2016-01-07 15:00'), + RESTRICTION_STATES.FROZEN), + Restriction(1, str_to_ts('2016-01-07 15:30'), + RESTRICTION_STATES.ALLOWED), + ]) + + for info in minutes_to_check: + bar_data = self.create_bardata( + simulation_dt_func=lambda: info[0], + restrictions=rlm, + ) + self.assertEqual(bar_data.can_trade(self.ASSET1), info[1]) + + +class TestMinuteBarDataMultipleExchanges(WithCreateBarData, + WithBarDataChecks, + ZiplineTestCase): + + START_DATE = pd.Timestamp('2016-01-05', tz='UTC') + END_DATE = ASSET_FINDER_EQUITY_END_DATE = pd.Timestamp( + '2016-01-07', + tz='UTC', + ) + + ASSET_FINDER_EQUITY_SIDS = [1] + + @classmethod + def make_equity_minute_bar_data(cls): + # asset1 has trades every minute + yield 1, create_minute_df_for_asset( + cls.trading_calendar, + cls.equity_minute_bar_days[0], + cls.equity_minute_bar_days[-1], + ) + + @classmethod + def make_futures_info(cls): + return pd.DataFrame.from_dict( + { + 6: { + 'symbol': 'CLG06', + 'root_symbol': 'CL', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2005-12-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-01-20', tz='UTC'), + 'exchange': 'ICEUS', + }, + }, + orient='index', + ) + + @classmethod + def init_class_fixtures(cls): + super(TestMinuteBarDataMultipleExchanges, cls).init_class_fixtures() + cls.trading_calendar = get_calendar('CME') + def test_can_trade_multiple_exchange_closed(self): nyse_asset = self.asset_finder.retrieve_asset(1) ice_asset = self.asset_finder.retrieve_asset(6) @@ -639,70 +767,18 @@ class TestMinuteBarData(WithBarDataChecks, for info in minutes_to_check: # use the CME calendar, which covers 24 hours - bar_data = BarData(self.data_portal, lambda: info[0], "minute", - trading_calendar=get_calendar("CME")) + bar_data = self.create_bardata( + simulation_dt_func=lambda: info[0], + ) series = bar_data.can_trade([nyse_asset, ice_asset]) self.assertEqual(info[1], series.loc[nyse_asset]) self.assertEqual(info[2], series.loc[ice_asset]) - def test_is_stale_during_non_market_hours(self): - bar_data = BarData( - self.data_portal, - lambda: self.equity_minute_bar_days[1], - "minute", - self.trading_calendar - ) - with handle_non_market_minutes(bar_data): - self.assertTrue(bar_data.is_stale(self.HILARIOUSLY_ILLIQUID_ASSET)) - - def test_overnight_adjustments(self): - # verify there is a split for SPLIT_ASSET - splits = self.adjustment_reader.get_adjustments_for_sid( - "splits", - self.SPLIT_ASSET.sid - ) - - self.assertEqual(1, len(splits)) - split = splits[0] - self.assertEqual( - split[0], - pd.Timestamp("2016-01-06", tz='UTC') - ) - - # Current day is 1/06/16 - day = self.equity_daily_bar_days[1] - eight_fortyfive_am_eastern = \ - pd.Timestamp("{0}-{1}-{2} 8:45".format( - day.year, day.month, day.day), - tz='US/Eastern' - ) - - bar_data = BarData(self.data_portal, - lambda: eight_fortyfive_am_eastern, - "minute", - self.trading_calendar) - - expected = { - 'open': 391 / 2.0, - 'high': 392 / 2.0, - 'low': 389 / 2.0, - 'close': 390 / 2.0, - 'volume': 39000 * 2.0, - 'price': 390 / 2.0, - } - - with handle_non_market_minutes(bar_data): - for field in OHLCP + ['volume']: - value = bar_data.current(self.SPLIT_ASSET, field) - - # Assert the price is adjusted for the overnight split - self.assertEqual(value, expected[field]) - - -class TestDailyBarData(WithBarDataChecks, +class TestDailyBarData(WithCreateBarData, + WithBarDataChecks, WithDataPortal, ZiplineTestCase): START_DATE = pd.Timestamp('2016-01-05', tz='UTC') @@ -710,6 +786,7 @@ class TestDailyBarData(WithBarDataChecks, '2016-01-11', tz='UTC', ) + CREATE_BARDATA_DATA_FREQUENCY = 'daily' sids = ASSET_FINDER_EQUITY_SIDS = set(range(1, 9)) @@ -848,8 +925,9 @@ class TestDailyBarData(WithBarDataChecks, ) ) - bar_data = BarData(self.data_portal, lambda: minute, "daily", - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: minute, + ) self.check_internal_consistency(bar_data) self.assertFalse(bar_data.can_trade(self.ASSET1)) @@ -871,13 +949,10 @@ class TestDailyBarData(WithBarDataChecks, def test_semi_active_day(self): # on self.equity_daily_bar_days[0], only asset1 has data - bar_data = BarData( - self.data_portal, - lambda: self.get_last_minute_of_session( + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.get_last_minute_of_session( self.equity_daily_bar_days[0] ), - "daily", - self.trading_calendar ) self.check_internal_consistency(bar_data) @@ -909,13 +984,10 @@ class TestDailyBarData(WithBarDataChecks, ) def test_fully_active_day(self): - bar_data = BarData( - self.data_portal, - lambda: self.get_last_minute_of_session( + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.get_last_minute_of_session( self.equity_daily_bar_days[1] ), - "daily", - self.trading_calendar ) self.check_internal_consistency(bar_data) @@ -936,13 +1008,10 @@ class TestDailyBarData(WithBarDataChecks, ) def test_last_active_day(self): - bar_data = BarData( - self.data_portal, - lambda: self.get_last_minute_of_session( + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.get_last_minute_of_session( self.equity_daily_bar_days[-1] ), - "daily", - self.trading_calendar ) self.check_internal_consistency(bar_data) @@ -971,8 +1040,9 @@ class TestDailyBarData(WithBarDataChecks, def test_after_assets_dead(self): session = self.END_DATE - bar_data = BarData(self.data_portal, lambda: session, "daily", - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: session, + ) self.check_internal_consistency(bar_data) for asset in self.ASSETS: @@ -1022,21 +1092,15 @@ class TestDailyBarData(WithBarDataChecks, ) # ... but that's it's not applied when using spot value - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[0], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[0], ) self.assertEqual( liquid_day_0_price, bar_data.current(liquid_asset, "price") ) - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[1], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[1], ) self.assertEqual( liquid_day_1_price, @@ -1045,21 +1109,15 @@ class TestDailyBarData(WithBarDataChecks, # ... except when we have to forward fill across a day boundary # ILLIQUID_ASSET has no data on days 0 and 2, and a split on day 2 - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[1], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[1], ) self.assertEqual( illiquid_day_0_price, bar_data.current(illiquid_asset, "price") ) - bar_data = BarData( - self.data_portal, - lambda: self.equity_daily_bar_days[2], - "daily", - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.equity_daily_bar_days[2], ) # 3 (price from previous day) * 0.5 (split ratio) @@ -1067,3 +1125,29 @@ class TestDailyBarData(WithBarDataChecks, illiquid_day_1_price_adjusted, bar_data.current(illiquid_asset, "price") ) + + def test_can_trade_restricted(self): + """ + Test that can_trade will return False for a sid if it is restricted + on that dt + """ + + minutes_to_check = [ + (pd.Timestamp("2016-01-05", tz="UTC"), False), + (pd.Timestamp("2016-01-06", tz="UTC"), False), + (pd.Timestamp("2016-01-07", tz="UTC"), True), + ] + + rlm = HistoricalRestrictions([ + Restriction(1, str_to_ts('2016-01-05'), + RESTRICTION_STATES.FROZEN), + Restriction(1, str_to_ts('2016-01-07'), + RESTRICTION_STATES.ALLOWED), + ]) + + for info in minutes_to_check: + bar_data = self.create_bardata( + simulation_dt_func=lambda: info[0], + restrictions=rlm + ) + self.assertEqual(bar_data.can_trade(self.ASSET1), info[1]) diff --git a/zipline/_protocol.pyx b/zipline/_protocol.pyx index fd0c39b2..b15159fd 100644 --- a/zipline/_protocol.pyx +++ b/zipline/_protocol.pyx @@ -153,6 +153,9 @@ cdef class BarData: data_frequency : {'minute', 'daily'} The frequency of the bar data; i.e. whether the data is daily or minute bars + restrictions : zipline.finance.restrictions.Restrictions + Object that combines and returns restricted list information from + multiple sources universe_func : callable, optional Function which returns the current 'universe'. This is for backwards compatibility with older API concepts. @@ -160,17 +163,19 @@ cdef class BarData: cdef object data_portal cdef object simulation_dt_func cdef object data_frequency + cdef object restrictions cdef dict _views cdef object _universe_func cdef object _last_calculated_universe cdef object _universe_last_updated_at cdef bool _daily_mode cdef object _trading_calendar + cdef object _is_restricted cdef bool _adjust_minutes def __init__(self, data_portal, simulation_dt_func, data_frequency, - trading_calendar, universe_func=None): + trading_calendar, restrictions, universe_func=None): self.data_portal = data_portal self.simulation_dt_func = simulation_dt_func self.data_frequency = data_frequency @@ -185,6 +190,7 @@ cdef class BarData: self._adjust_minutes = False self._trading_calendar = trading_calendar + self._is_restricted = restrictions.is_restricted cdef _get_equity_price_view(self, asset): """ @@ -482,6 +488,9 @@ cdef class BarData: cdef object session_label cdef object dt_to_use_for_exchange_check, + if self._is_restricted(asset, adjusted_dt): + return False + session_label = self._trading_calendar.minute_to_session_label(dt) if not asset.is_alive_for_session(session_label): diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 86553d27..9b6df86c 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -76,11 +76,16 @@ from zipline.finance.execution import ( StopOrder, ) from zipline.finance.performance import PerformanceTracker +from zipline.finance.restrictions import Restrictions from zipline.finance.slippage import ( VolumeShareSlippage, SlippageModel ) from zipline.finance.cancel_policy import NeverCancel, CancelPolicy +from zipline.finance.restrictions import ( + NoopRestrictions, + StaticRestrictions +) from zipline.assets import Asset, Future from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.pipeline import Pipeline @@ -120,6 +125,7 @@ from zipline.utils.math_utils import ( round_if_near_integer ) from zipline.utils.preprocess import preprocess +from zipline.utils.security_list import SecurityList import zipline.protocol from zipline.sources.requests_csv import PandasRequestsCSV @@ -418,6 +424,8 @@ class TradingAlgorithm(object): # A dictionary of the actual capital change deltas, keyed by timestamp self.capital_change_deltas = {} + self.restrictions = NoopRestrictions() + def init_engine(self, get_loader): """ Construct and store a PipelineEngine from loader. @@ -564,6 +572,7 @@ class TradingAlgorithm(object): self.data_portal, self._create_clock(), self._create_benchmark_source(), + self.restrictions, universe_func=self._calculate_universe ) @@ -2083,7 +2092,8 @@ class TradingAlgorithm(object): def set_max_position_size(self, asset=None, max_shares=None, - max_notional=None): + max_notional=None, + on_error='fail'): """Set a limit on the number of shares and/or dollar value held for the given sid. Limits are treated as absolute values and are enforced at the time that the algo attempts to place an order for sid. This means @@ -2107,14 +2117,16 @@ class TradingAlgorithm(object): """ control = MaxPositionSize(asset=asset, max_shares=max_shares, - max_notional=max_notional) + max_notional=max_notional, + on_error=on_error) self.register_trading_control(control) @api_method def set_max_order_size(self, asset=None, max_shares=None, - max_notional=None): + max_notional=None, + on_error='fail'): """Set a limit on the number of shares and/or dollar value of any single order placed for sid. Limits are treated as absolute values and are enforced at the time that the algo attempts to place an order for sid. @@ -2134,11 +2146,12 @@ class TradingAlgorithm(object): """ control = MaxOrderSize(asset=asset, max_shares=max_shares, - max_notional=max_notional) + max_notional=max_notional, + on_error=on_error) self.register_trading_control(control) @api_method - def set_max_order_count(self, max_count): + def set_max_order_count(self, max_count, on_error='fail'): """Set a limit on the number of orders that can be placed in a single day. @@ -2147,27 +2160,32 @@ class TradingAlgorithm(object): max_count : int The maximum number of orders that can be placed on any single day. """ - control = MaxOrderCount(max_count) + control = MaxOrderCount(on_error, max_count) self.register_trading_control(control) @api_method - def set_do_not_order_list(self, restricted_list): + def set_do_not_order_list(self, restricted_list, on_error='fail'): """Set a restriction on which assets can be ordered. Parameters ---------- - restricted_list : container[Asset] + restricted_list : container[Asset], SecurityList The assets that cannot be ordered. """ - control = RestrictedListOrder(restricted_list) + + if isinstance(restricted_list, (list, tuple, set)): + restricted_list = StaticRestrictions(restricted_list) + + control = RestrictedListOrder(on_error, restricted_list) self.register_trading_control(control) + self.restrictions = restricted_list @api_method - def set_long_only(self): + def set_long_only(self, on_error='fail'): """Set a rule specifying that this algorithm cannot take short positions. """ - self.register_trading_control(LongOnly()) + self.register_trading_control(LongOnly(on_error)) ############## # Pipeline API diff --git a/zipline/api.py b/zipline/api.py index ef3a75f8..426b5f26 100644 --- a/zipline/api.py +++ b/zipline/api.py @@ -16,6 +16,12 @@ # Note that part of the API is implemented in TradingAlgorithm as # methods (e.g. order). These are added to this namespace via the # decorator ``api_method`` inside of algorithm.py. +from .finance.restrictions import ( + Restriction, + StaticRestrictions, + HistoricalRestrictions, + RESTRICTION_STATES, +) from .finance import commission, execution, slippage, cancel_policy from .finance.cancel_policy import ( NeverCancel, @@ -36,6 +42,10 @@ __all__ = [ 'FixedSlippage', 'NeverCancel', 'VolumeShareSlippage', + 'Restriction', + 'StaticRestrictions', + 'HistoricalRestrictions', + 'RESTRICTION_STATES', 'cancel_policy', 'commission', 'date_rules', diff --git a/zipline/finance/controls.py b/zipline/finance/controls.py index b1c46de8..95d61c89 100644 --- a/zipline/finance/controls.py +++ b/zipline/finance/controls.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc +import logbook import pandas as pd @@ -23,6 +24,8 @@ from zipline.errors import ( TradingControlViolation, ) +log = logbook.Logger('TradingControl') + class TradingControl(with_metaclass(abc.ABCMeta)): """ @@ -30,11 +33,12 @@ class TradingControl(with_metaclass(abc.ABCMeta)): algorithm. """ - def __init__(self, **kwargs): + def __init__(self, on_error, **kwargs): """ Track any arguments that should be printed in the error message generated by self.fail. """ + self.on_error = on_error self.__fail_args = kwargs @abc.abstractmethod @@ -57,23 +61,36 @@ class TradingControl(with_metaclass(abc.ABCMeta)): """ raise NotImplementedError - def fail(self, asset, amount, datetime, metadata=None): - """ - Raise a TradingControlViolation with information about the failure. - - If dynamic information should be displayed as well, pass it in via - `metadata`. - """ + def _constraint_msg(self, metadata): constraint = repr(self) if metadata: constraint = "{constraint} (Metadata: {metadata})".format( constraint=constraint, metadata=metadata ) - raise TradingControlViolation(asset=asset, - amount=amount, - datetime=datetime, - constraint=constraint) + return constraint + + def handle_violation(self, asset, amount, datetime, metadata=None): + """ + Handle a TradingControlViolation, either by raising or logging and + error with information about the failure. + + If dynamic information should be displayed as well, pass it in via + `metadata`. + """ + constraint = self._constraint_msg(metadata) + + if self.on_error == 'fail': + raise TradingControlViolation( + asset=asset, + amount=amount, + datetime=datetime, + constraint=constraint) + elif self.on_error == 'log': + log.error("Order for {amount} shares of {asset} at {dt} " + "violates trading constraint {constraint}", + amount=amount, asset=asset, dt=datetime, + constraint=constraint) def __repr__(self): return "{name}({attrs})".format(name=self.__class__.__name__, @@ -86,9 +103,9 @@ class MaxOrderCount(TradingControl): placed in a given trading day. """ - def __init__(self, max_count): + def __init__(self, on_error, max_count): - super(MaxOrderCount, self).__init__(max_count=max_count) + super(MaxOrderCount, self).__init__(on_error, max_count=max_count) self.orders_placed = 0 self.max_count = max_count self.current_date = None @@ -96,9 +113,9 @@ class MaxOrderCount(TradingControl): def validate(self, asset, amount, - _portfolio, + portfolio, algo_datetime, - _algo_current_data): + algo_current_data): """ Fail if we've already placed self.max_count orders today. """ @@ -110,7 +127,7 @@ class MaxOrderCount(TradingControl): self.current_date = algo_date if self.orders_placed >= self.max_count: - self.fail(asset, amount, algo_datetime) + self.handle_violation(asset, amount, algo_datetime) self.orders_placed += 1 @@ -120,25 +137,25 @@ class RestrictedListOrder(TradingControl): Parameters ---------- - restricted_list : container[Asset] - The assets that cannot be ordered. + restrictions : zipline.finance.restrictions.Restrictions + Object representing restrictions of a group of assets. """ - def __init__(self, restricted_list): - super(RestrictedListOrder, self).__init__() - self.restricted_list = restricted_list + def __init__(self, on_error, restrictions): + super(RestrictedListOrder, self).__init__(on_error) + self.restrictions = restrictions def validate(self, asset, amount, - _portfolio, - _algo_datetime, - _algo_current_data): + portfolio, + algo_datetime, + algo_current_data): """ Fail if the asset is in the restricted_list. """ - if asset in self.restricted_list: - self.fail(asset, amount, _algo_datetime) + if self.restrictions.is_restricted(asset, algo_datetime): + self.handle_violation(asset, amount, algo_datetime) class MaxOrderSize(TradingControl): @@ -148,8 +165,10 @@ class MaxOrderSize(TradingControl): value. """ - def __init__(self, asset=None, max_shares=None, max_notional=None): - super(MaxOrderSize, self).__init__(asset=asset, + def __init__(self, on_error, asset=None, max_shares=None, + max_notional=None): + super(MaxOrderSize, self).__init__(on_error, + asset=asset, max_shares=max_shares, max_notional=max_notional) self.asset = asset @@ -175,7 +194,7 @@ class MaxOrderSize(TradingControl): asset, amount, portfolio, - _algo_datetime, + algo_datetime, algo_current_data): """ Fail if the magnitude of the given order exceeds either self.max_shares @@ -186,7 +205,7 @@ class MaxOrderSize(TradingControl): return if self.max_shares is not None and abs(amount) > self.max_shares: - self.fail(asset, amount, _algo_datetime) + self.handle_violation(asset, amount, algo_datetime) current_asset_price = algo_current_data.current(asset, "price") order_value = amount * current_asset_price @@ -195,7 +214,7 @@ class MaxOrderSize(TradingControl): abs(order_value) > self.max_notional) if too_much_value: - self.fail(asset, amount, _algo_datetime) + self.handle_violation(asset, amount, algo_datetime) class MaxPositionSize(TradingControl): @@ -204,8 +223,10 @@ class MaxPositionSize(TradingControl): be held by an algo for a given asset. """ - def __init__(self, asset=None, max_shares=None, max_notional=None): - super(MaxPositionSize, self).__init__(asset=asset, + def __init__(self, on_error, asset=None, max_shares=None, + max_notional=None): + super(MaxPositionSize, self).__init__(on_error, + asset=asset, max_shares=max_shares, max_notional=max_notional) self.asset = asset @@ -248,7 +269,7 @@ class MaxPositionSize(TradingControl): too_many_shares = (self.max_shares is not None and abs(shares_post_order) > self.max_shares) if too_many_shares: - self.fail(asset, amount, algo_datetime) + self.handle_violation(asset, amount, algo_datetime) current_price = algo_current_data.current(asset, "price") value_post_order = shares_post_order * current_price @@ -257,7 +278,7 @@ class MaxPositionSize(TradingControl): abs(value_post_order) > self.max_notional) if too_much_value: - self.fail(asset, amount, algo_datetime) + self.handle_violation(asset, amount, algo_datetime) class LongOnly(TradingControl): @@ -265,18 +286,21 @@ class LongOnly(TradingControl): TradingControl representing a prohibition against holding short positions. """ + def __init__(self, on_error): + super(LongOnly, self).__init__(on_error) + def validate(self, asset, amount, portfolio, - _algo_datetime, - _algo_current_data): + algo_datetime, + algo_current_data): """ Fail if we would hold negative shares of asset after completing this order. """ if portfolio.positions[asset].amount + amount < 0: - self.fail(asset, amount, _algo_datetime) + self.handle_violation(asset, amount, algo_datetime) class AssetDateBounds(TradingControl): @@ -285,6 +309,9 @@ class AssetDateBounds(TradingControl): its start_date, or after its end_date. """ + def __init__(self, on_error): + super(AssetDateBounds, self).__init__(on_error) + def validate(self, asset, amount, @@ -308,7 +335,8 @@ class AssetDateBounds(TradingControl): metadata = { 'asset_start_date': normalized_start } - self.fail(asset, amount, algo_datetime, metadata=metadata) + self.handle_violation( + asset, amount, algo_datetime, metadata=metadata) # Fail if the algo has passed this Asset's end_date if asset.end_date: normalized_end = pd.Timestamp(asset.end_date).normalize() @@ -316,7 +344,8 @@ class AssetDateBounds(TradingControl): metadata = { 'asset_end_date': normalized_end } - self.fail(asset, amount, algo_datetime, metadata=metadata) + self.handle_violation( + asset, amount, algo_datetime, metadata=metadata) class AccountControl(with_metaclass(abc.ABCMeta)): diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 2ab816d7..1d5b372f 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -38,7 +38,7 @@ class AlgorithmSimulator(object): } def __init__(self, algo, sim_params, data_portal, clock, benchmark_source, - universe_func): + restrictions, universe_func): # ============== # Simulation @@ -47,6 +47,7 @@ class AlgorithmSimulator(object): self.sim_params = sim_params self.env = algo.trading_environment self.data_portal = data_portal + self.restrictions = restrictions # ============== # Algo Setup @@ -89,6 +90,7 @@ class AlgorithmSimulator(object): simulation_dt_func=self.get_simulation_dt, data_frequency=self.sim_params.data_frequency, trading_calendar=self.algo.trading_calendar, + restrictions=self.restrictions, universe_func=universe_func ) diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index d857245a..052dd821 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -505,9 +505,9 @@ class SetMaxOrderSizeAlgorithm(TradingAlgorithm): class SetDoNotOrderListAlgorithm(TradingAlgorithm): - def initialize(self, sid=None, restricted_list=None): + def initialize(self, sid=None, restricted_list=None, on_error='fail'): self.order_count = 0 - self.set_do_not_order_list(restricted_list) + self.set_do_not_order_list(restricted_list, on_error) class SetMaxOrderCountAlgorithm(TradingAlgorithm): @@ -529,7 +529,7 @@ class SetAssetDateBoundsAlgorithm(TradingAlgorithm): AssetDateBounds() trading control in place. """ def initialize(self): - self.register_trading_control(AssetDateBounds()) + self.register_trading_control(AssetDateBounds(on_error='fail')) def handle_data(algo, data): algo.order(algo.sid(999), 1) From e465f64f912118258a78a5dbeb0e89cfcf109bed Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Wed, 14 Sep 2016 14:23:40 -0400 Subject: [PATCH 03/15] MAINT: Create SecurityListRestrictions that takes a SecurityList The SecurityList implements a non-exposed method `current_securities(dt)` which SecurityListRestrictions calls to determine if an asset is restricted. Deprecate the `__iter__` and `__contains__` methods of security lists in favor of `current_securities(dt)` --- tests/test_restrictions.py | 43 +++++++++++++++++++++++++++++++++ tests/test_security_list.py | 27 +++++++++++++++------ zipline/algorithm.py | 5 +++- zipline/api.py | 2 ++ zipline/finance/restrictions.py | 23 ++++++++++++++++++ zipline/utils/security_list.py | 25 +++++++++++++------ 6 files changed, 110 insertions(+), 15 deletions(-) diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py index c1817b68..74327ce9 100644 --- a/tests/test_restrictions.py +++ b/tests/test_restrictions.py @@ -9,6 +9,7 @@ from zipline.finance.restrictions import ( Restriction, HistoricalRestrictions, StaticRestrictions, + SecurityListRestrictions, NoopRestrictions, ) @@ -212,6 +213,48 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): assert_vectorized_results([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_vectorized_results = partial(self.assert_vectorized_results, 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_vectorized_results( + [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_vectorized_results([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_vectorized_results([True, True, True], str_to_ts('2011-01-05')) + def test_noop_restrictions(self): """ Test single- and multi-asset queries on no-op restrictions diff --git a/tests/test_security_list.py b/tests/test_security_list.py index f0609142..6dd975c9 100644 --- a/tests/test_security_list.py +++ b/tests/test_security_list.py @@ -36,7 +36,8 @@ class RestrictedAlgoWithCheck(TradingAlgorithm): def handle_data(self, data): if not self.order_count: if self.sid not in \ - self.rl.leveraged_etf_list: + self.rl.leveraged_etf_list.\ + current_securities(self.get_datetime()): self.order(self.sid, 100) self.order_count += 1 @@ -62,7 +63,8 @@ class IterateRLAlgo(TradingAlgorithm): self.found = False def handle_data(self, data): - for stock in self.rl.leveraged_etf_list: + for stock in self.rl.leveraged_etf_list.\ + current_securities(self.get_datetime()): if stock == self.sid: self.found = True @@ -151,7 +153,8 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): for symbol in ["BZQ", "URTY", "JFT"]] ] for sid in should_exist: - self.assertIn(sid, rl.leveraged_etf_list) + self.assertIn( + sid, rl.leveraged_etf_list.current_securities(get_datetime())) # assert that a sample of allowed stocks are not in restricted shouldnt_exist = [ @@ -162,7 +165,8 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): for symbol in ["AAPL", "GOOG"]] ] for sid in shouldnt_exist: - self.assertNotIn(sid, rl.leveraged_etf_list) + self.assertNotIn( + sid, rl.leveraged_etf_list.current_securities(get_datetime())) def test_security_add(self): def get_datetime(): @@ -178,15 +182,24 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): ) for symbol in ["AAPL", "GOOG", "BZQ", "URTY"]] ] for sid in should_exist: - self.assertIn(sid, rl.leveraged_etf_list) + self.assertIn( + sid, + rl.leveraged_etf_list.current_securities(get_datetime()) + ) def test_security_add_delete(self): with security_list_copy(): def get_datetime(): return pd.Timestamp("2015-01-27", tz='UTC') rl = SecurityListSet(get_datetime, self.env.asset_finder) - self.assertNotIn("BZQ", rl.leveraged_etf_list) - self.assertNotIn("URTY", rl.leveraged_etf_list) + self.assertNotIn( + "BZQ", + rl.leveraged_etf_list.current_securities(get_datetime()) + ) + self.assertNotIn( + "URTY", + rl.leveraged_etf_list.current_securities(get_datetime()) + ) def test_algo_without_rl_violation_via_check(self): algo = RestrictedAlgoWithCheck(symbol='BZQ', diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 9b6df86c..d5dc429f 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -84,7 +84,8 @@ from zipline.finance.slippage import ( from zipline.finance.cancel_policy import NeverCancel, CancelPolicy from zipline.finance.restrictions import ( NoopRestrictions, - StaticRestrictions + StaticRestrictions, + SecurityListRestrictions, ) from zipline.assets import Asset, Future from zipline.gens.tradesimulation import AlgorithmSimulator @@ -2175,6 +2176,8 @@ class TradingAlgorithm(object): if isinstance(restricted_list, (list, tuple, set)): restricted_list = StaticRestrictions(restricted_list) + elif isinstance(restricted_list, SecurityList): + restricted_list = SecurityListRestrictions(restricted_list) control = RestrictedListOrder(on_error, restricted_list) self.register_trading_control(control) diff --git a/zipline/api.py b/zipline/api.py index 426b5f26..f3e345ac 100644 --- a/zipline/api.py +++ b/zipline/api.py @@ -21,6 +21,7 @@ from .finance.restrictions import ( StaticRestrictions, HistoricalRestrictions, RESTRICTION_STATES, + SecurityListRestrictions, ) from .finance import commission, execution, slippage, cancel_policy from .finance.cancel_policy import ( @@ -46,6 +47,7 @@ __all__ = [ 'StaticRestrictions', 'HistoricalRestrictions', 'RESTRICTION_STATES', + 'SecurityListRestrictions', 'cancel_policy', 'commission', 'date_rules', diff --git a/zipline/finance/restrictions.py b/zipline/finance/restrictions.py index e977d2d8..52d3b150 100644 --- a/zipline/finance/restrictions.py +++ b/zipline/finance/restrictions.py @@ -127,3 +127,26 @@ class HistoricalRestrictions(Restrictions): break state = r.state return state == RESTRICTION_STATES.FROZEN + + +class SecurityListRestrictions(Restrictions): + """ + Restrictions based on a security list + + Parameters + ---------- + restrictions : zipline.utils.security_list.SecurityList + The restrictions defined by a SecurityList + """ + + def __init__(self, security_list_by_dt): + self.current_securities = security_list_by_dt.current_securities + + def is_restricted(self, assets, dt): + securities_in_list = self.current_securities(dt) + if isinstance(assets, Asset): + return assets in securities_in_list + return pd.Series( + index=pd.Index(assets), + data=vectorized_is_element(assets, securities_in_list) + ) diff --git a/zipline/utils/security_list.py b/zipline/utils/security_list.py index 50368b59..8532d3fb 100644 --- a/zipline/utils/security_list.py +++ b/zipline/utils/security_list.py @@ -1,3 +1,4 @@ +import warnings from datetime import datetime from os import listdir import os.path @@ -7,6 +8,7 @@ import pytz import zipline from zipline.errors import SymbolNotFound +from zipline.zipline_warnings import ZiplineDeprecationWarning DATE_FORMAT = "%Y%m%d" @@ -38,17 +40,26 @@ class SecurityList(object): return knowledge_dates def __iter__(self): - return iter(self.restricted_list) + warnings.warn( + 'Iterating over security_lists is deprecated. Use ' + '`for sid in .current_securities(dt)` instead.', + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + return iter(self.current_securities(self.current_date())) def __contains__(self, item): - return item in self.restricted_list + warnings.warn( + 'Evaluating inclusion in security_lists is deprecated. Use ' + '`sid in .current_securities(dt)` instead.', + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + return item in self.current_securities(self.current_date()) - @property - def restricted_list(self): - - cd = self.current_date() + def current_securities(self, dt): for kd in self._knowledge_dates: - if cd < kd: + if dt < kd: break if kd in self._cache: self._current_set = self._cache[kd] From 5e276d0e72f2ca560480138e69f595628ab2390c Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Wed, 14 Sep 2016 11:29:52 -0400 Subject: [PATCH 04/15] TEST: Modify tests for extra BarData parameter Introducing a WithCreateBarData fixture which allows for the creation of a BarData using only the `simulation_dt_func` and `restrictions` params. Assumes that each suite uses the same `data_portal`, `data_frequency` and `trading_calendar` --- tests/finance/test_slippage.py | 783 +++++++++++++++++---------------- tests/test_api_shim.py | 13 +- tests/test_blotter.py | 24 +- tests/test_finance.py | 10 +- tests/test_history.py | 86 ++-- tests/test_tradesimulation.py | 2 + zipline/testing/fixtures.py | 16 + 7 files changed, 502 insertions(+), 432 deletions(-) diff --git a/tests/finance/test_slippage.py b/tests/finance/test_slippage.py index 2c6ae008..17846b94 100644 --- a/tests/finance/test_slippage.py +++ b/tests/finance/test_slippage.py @@ -27,20 +27,25 @@ from pandas.tslib import normalize_date from zipline.finance.slippage import VolumeShareSlippage -from zipline.protocol import DATASOURCE_TYPE +from zipline.protocol import DATASOURCE_TYPE, BarData from zipline.finance.blotter import Order - +from zipline.finance.restrictions import NoopRestrictions from zipline.data.data_portal import DataPortal -from zipline.protocol import BarData from zipline.testing import tmp_bcolz_equity_minute_bar_reader from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, WithSimParams, + WithTradingEnvironment, ZiplineTestCase, ) +from zipline.utils.classproperty import classproperty -class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): +class SlippageTestCase(WithCreateBarData, + WithSimParams, + WithDataPortal, + ZiplineTestCase): START_DATE = pd.Timestamp('2006-01-05 14:31', tz='utc') END_DATE = pd.Timestamp('2006-01-05 14:36', tz='utc') SIM_PARAMS_CAPITAL_BASE = 1.0e5 @@ -56,6 +61,10 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): freq='1min' ) + @classproperty + def CREATE_BARDATA_DATA_FREQUENCY(cls): + return cls.sim_params.data_frequency + @classmethod def make_equity_minute_bar_data(cls): yield 133, pd.DataFrame( @@ -74,97 +83,6 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): super(SlippageTestCase, cls).init_class_fixtures() cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133) - def test_volume_share_slippage(self): - assets = ( - (133, pd.DataFrame( - { - 'open': [3.00], - 'high': [3.15], - 'low': [2.85], - 'close': [3.00], - 'volume': [200], - }, - index=[self.minutes[0]], - )), - ) - days = pd.date_range( - start=normalize_date(self.minutes[0]), - end=normalize_date(self.minutes[-1]) - ) - with tmp_bcolz_equity_minute_bar_reader(self.trading_calendar, days, assets) \ - as reader: - data_portal = DataPortal( - self.env.asset_finder, self.trading_calendar, - first_trading_day=reader.first_trading_day, - equity_minute_reader=reader, - ) - - slippage_model = VolumeShareSlippage() - - open_orders = [ - Order( - dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - amount=100, - filled=0, - sid=self.ASSET133 - ) - ] - - bar_data = BarData(data_portal, - lambda: self.minutes[0], - 'minute', - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 1) - _, txn = orders_txns[0] - - expected_txn = { - 'price': float(3.0001875), - 'dt': datetime.datetime( - 2006, 1, 5, 14, 31, tzinfo=pytz.utc), - 'amount': int(5), - 'sid': int(133), - 'commission': None, - 'type': DATASOURCE_TYPE.TRANSACTION, - 'order_id': open_orders[0].id - } - - self.assertIsNotNone(txn) - - # TODO: Make expected_txn an Transaction object and ensure there - # is a __eq__ for that class. - self.assertEquals(expected_txn, txn.__dict__) - - open_orders = [ - Order( - dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - amount=100, - filled=0, - sid=self.ASSET133 - ) - ] - - # Set bar_data to be a minute ahead of last trade. - # Volume share slippage should not execute when there is no trade. - bar_data = BarData(data_portal, - lambda: self.minutes[1], - 'minute', - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - def test_orders_limit(self): slippage_model = VolumeShareSlippage() slippage_model.data_portal = self.data_portal @@ -179,10 +97,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -202,10 +119,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -225,10 +141,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.6}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -265,10 +180,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -288,10 +202,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.5}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -311,10 +224,9 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): 'limit': 3.4}) ] - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) orders_txns = list(slippage_model.simulate( bar_data, @@ -338,6 +250,376 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): for key, value in expected_txn.items(): self.assertEquals(value, txn[key]) + def test_orders_stop_limit(self): + slippage_model = VolumeShareSlippage() + slippage_model.data_portal = self.data_portal + + # long, does not trade + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': 100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 4.0, + 'limit': 3.0}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[2], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # long, does not trade - impacted price worse than limit price + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': 100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 4.0, + 'limit': 3.5}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[2], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # long, does trade + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': 100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 4.0, + 'limit': 3.6}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[2], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[3], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 1) + _, txn = orders_txns[0] + + expected_txn = { + 'price': float(3.50021875), + 'dt': datetime.datetime( + 2006, 1, 5, 14, 34, tzinfo=pytz.utc), + 'amount': int(50), + 'sid': int(133) + } + + for key, value in expected_txn.items(): + self.assertEquals(value, txn[key]) + + # short, does not trade + + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': -100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 3.0, + 'limit': 4.0}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # short, does not trade - impacted price worse than limit price + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': -100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 3.0, + 'limit': 3.5}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + # short, does trade + open_orders = [ + Order(**{ + 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + 'amount': -100, + 'filled': 0, + 'sid': self.ASSET133, + 'stop': 3.0, + 'limit': 3.4}) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 1) + _, txn = orders_txns[0] + + expected_txn = { + 'price': float(3.49978125), + 'dt': datetime.datetime( + 2006, 1, 5, 14, 32, tzinfo=pytz.utc), + 'amount': int(-50), + 'sid': int(133) + } + + for key, value in expected_txn.items(): + self.assertEquals(value, txn[key]) + + +class VolumeShareSlippageTestCase(WithCreateBarData, + WithSimParams, + WithDataPortal, + ZiplineTestCase): + + START_DATE = pd.Timestamp('2006-01-05 14:31', tz='utc') + END_DATE = pd.Timestamp('2006-01-05 14:36', tz='utc') + SIM_PARAMS_CAPITAL_BASE = 1.0e5 + SIM_PARAMS_DATA_FREQUENCY = 'minute' + SIM_PARAMS_EMISSION_RATE = 'daily' + + ASSET_FINDER_EQUITY_SIDS = (133,) + ASSET_FINDER_EQUITY_START_DATE = pd.Timestamp('2006-01-05', tz='utc') + ASSET_FINDER_EQUITY_END_DATE = pd.Timestamp('2006-01-07', tz='utc') + minutes = pd.DatetimeIndex( + start=START_DATE, + end=END_DATE - pd.Timedelta('1 minute'), + freq='1min' + ) + + @classproperty + def CREATE_BARDATA_DATA_FREQUENCY(cls): + return cls.sim_params.data_frequency + + @classmethod + def make_equity_minute_bar_data(cls): + yield 133, pd.DataFrame( + { + 'open': [3.00], + 'high': [3.15], + 'low': [2.85], + 'close': [3.00], + 'volume': [200], + }, + index=[cls.minutes[0]], + ) + + @classmethod + def init_class_fixtures(cls): + super(VolumeShareSlippageTestCase, cls).init_class_fixtures() + cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133) + + def test_volume_share_slippage(self): + + slippage_model = VolumeShareSlippage() + + open_orders = [ + Order( + dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + amount=100, + filled=0, + sid=self.ASSET133 + ) + ] + + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[0], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 1) + _, txn = orders_txns[0] + + expected_txn = { + 'price': float(3.0001875), + 'dt': datetime.datetime( + 2006, 1, 5, 14, 31, tzinfo=pytz.utc), + 'amount': int(5), + 'sid': int(133), + 'commission': None, + 'type': DATASOURCE_TYPE.TRANSACTION, + 'order_id': open_orders[0].id + } + + self.assertIsNotNone(txn) + + # TODO: Make expected_txn an Transaction object and ensure there + # is a __eq__ for that class. + self.assertEquals(expected_txn, txn.__dict__) + + open_orders = [ + Order( + dt=datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), + amount=100, + filled=0, + sid=self.ASSET133 + ) + ] + + # Set bar_data to be a minute ahead of last trade. + # Volume share slippage should not execute when there is no trade. + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.minutes[1], + ) + + orders_txns = list(slippage_model.simulate( + bar_data, + self.ASSET133, + open_orders, + )) + + self.assertEquals(len(orders_txns), 0) + + +class OrdersStopTestCase(WithSimParams, + WithTradingEnvironment, + ZiplineTestCase): + + START_DATE = pd.Timestamp('2006-01-05 14:31', tz='utc') + END_DATE = pd.Timestamp('2006-01-05 14:36', tz='utc') + SIM_PARAMS_CAPITAL_BASE = 1.0e5 + SIM_PARAMS_DATA_FREQUENCY = 'minute' + SIM_PARAMS_EMISSION_RATE = 'daily' + ASSET_FINDER_EQUITY_SIDS = (133,) + minutes = pd.DatetimeIndex( + start=START_DATE, + end=END_DATE - pd.Timedelta('1 minute'), + freq='1min' + ) + + @classmethod + def init_class_fixtures(cls): + super(OrdersStopTestCase, cls).init_class_fixtures() + cls.ASSET133 = cls.env.asset_finder.retrieve_asset(133) + STOP_ORDER_CASES = { # Stop orders can be long/short and have their price greater or # less than the stop. @@ -501,10 +783,14 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): try: dt = pd.Timestamp('2006-01-05 14:31', tz='UTC') - bar_data = BarData(data_portal, - lambda: dt, - 'minute', - self.trading_calendar) + bar_data = BarData( + data_portal, + lambda: dt, + self.sim_params.data_frequency, + self.trading_calendar, + NoopRestrictions(), + ) + _, txn = next(slippage_model.simulate( bar_data, self.ASSET133, @@ -520,254 +806,3 @@ class SlippageTestCase(WithSimParams, WithDataPortal, ZiplineTestCase): for key, value in expected['transaction'].items(): self.assertEquals(value, txn[key]) - - def test_orders_stop_limit(self): - slippage_model = VolumeShareSlippage() - slippage_model.data_portal = self.data_portal - - # long, does not trade - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': 100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 4.0, - 'limit': 3.0}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[2], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # long, does not trade - impacted price worse than limit price - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': 100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 4.0, - 'limit': 3.5}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[2], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # long, does trade - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': 100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 4.0, - 'limit': 3.6}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[2], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[3], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 1) - _, txn = orders_txns[0] - - expected_txn = { - 'price': float(3.50021875), - 'dt': datetime.datetime( - 2006, 1, 5, 14, 34, tzinfo=pytz.utc), - 'amount': int(50), - 'sid': int(133) - } - - for key, value in expected_txn.items(): - self.assertEquals(value, txn[key]) - - # short, does not trade - - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': -100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 3.0, - 'limit': 4.0}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # short, does not trade - impacted price worse than limit price - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': -100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 3.0, - 'limit': 3.5}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - # short, does trade - open_orders = [ - Order(**{ - 'dt': datetime.datetime(2006, 1, 5, 14, 30, tzinfo=pytz.utc), - 'amount': -100, - 'filled': 0, - 'sid': self.ASSET133, - 'stop': 3.0, - 'limit': 3.4}) - ] - - bar_data = BarData(self.data_portal, - lambda: self.minutes[0], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 0) - - bar_data = BarData(self.data_portal, - lambda: self.minutes[1], - self.sim_params.data_frequency, - self.trading_calendar) - - orders_txns = list(slippage_model.simulate( - bar_data, - self.ASSET133, - open_orders, - )) - - self.assertEquals(len(orders_txns), 1) - _, txn = orders_txns[0] - - expected_txn = { - 'price': float(3.49978125), - 'dt': datetime.datetime( - 2006, 1, 5, 14, 32, tzinfo=pytz.utc), - 'amount': int(-50), - 'sid': int(133) - } - - for key, value in expected_txn.items(): - self.assertEquals(value, txn[key]) diff --git a/tests/test_api_shim.py b/tests/test_api_shim.py index 3aa64b8d..f79416d2 100644 --- a/tests/test_api_shim.py +++ b/tests/test_api_shim.py @@ -7,7 +7,6 @@ from pandas.core.common import PerformanceWarning from zipline import TradingAlgorithm from zipline.finance.trading import SimulationParameters -from zipline.protocol import BarData from zipline.testing import ( MockDailyBarReader, create_daily_df_for_asset, @@ -15,6 +14,7 @@ from zipline.testing import ( str_to_seconds, ) from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, WithSimParams, ZiplineTestCase, @@ -114,7 +114,11 @@ def handle_data(context, data): """ -class TestAPIShim(WithDataPortal, WithSimParams, ZiplineTestCase): +class TestAPIShim(WithCreateBarData, + WithDataPortal, + WithSimParams, + ZiplineTestCase, + ): START_DATE = pd.Timestamp("2016-01-05", tz='UTC') END_DATE = pd.Timestamp("2016-01-28", tz='UTC') SIM_PARAMS_DATA_FREQUENCY = 'minute' @@ -186,11 +190,8 @@ class TestAPIShim(WithDataPortal, WithSimParams, ZiplineTestCase): test_end_minute = self.trading_calendar.minutes_for_session( self.sim_params.sessions[0] )[-1] - bar_data = BarData( - self.data_portal, + bar_data = self.create_bardata( lambda: test_end_minute, - "minute", - self.trading_calendar ) ohlcvp_fields = [ "open", diff --git a/tests/test_blotter.py b/tests/test_blotter.py index 89c251db..fa5d0cf3 100644 --- a/tests/test_blotter.py +++ b/tests/test_blotter.py @@ -31,8 +31,9 @@ from zipline.finance.slippage import ( DEFAULT_VOLUME_SLIPPAGE_BAR_LIMIT, FixedSlippage, ) -from zipline.protocol import BarData +from zipline.utils.classproperty import classproperty from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, WithLogger, WithSimParams, @@ -40,7 +41,8 @@ from zipline.testing.fixtures import ( ) -class BlotterTestCase(WithLogger, +class BlotterTestCase(WithCreateBarData, + WithLogger, WithDataPortal, WithSimParams, ZiplineTestCase): @@ -71,6 +73,10 @@ class BlotterTestCase(WithLogger, index=cls.sim_params.sessions, ) + @classproperty + def CREATE_BARDATA_DATA_FREQUENCY(cls): + return cls.sim_params.data_frequency + @parameterized.expand([(MarketOrder(), None, None), (LimitOrder(10), 10, None), (StopOrder(10), None, 10), @@ -219,11 +225,8 @@ class BlotterTestCase(WithLogger, filled_id = blotter.order(asset_24, 100, MarketOrder()) filled_order = None blotter.current_dt = self.sim_params.sessions[-1] - bar_data = BarData( - self.data_portal, - lambda: self.sim_params.sessions[-1], - self.sim_params.data_frequency, - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: self.sim_params.sessions[-1], ) txns, _, closed_orders = blotter.get_transactions(bar_data) for txn in txns: @@ -295,11 +298,8 @@ class BlotterTestCase(WithLogger, filled_order = None blotter.current_dt = dt - bar_data = BarData( - self.data_portal, - lambda: dt, - self.sim_params.data_frequency, - self.trading_calendar + bar_data = self.create_bardata( + simulation_dt_func=lambda: dt, ) txns, _, _ = blotter.get_transactions(bar_data) for txn in txns: diff --git a/tests/test_finance.py b/tests/test_finance.py index 4b374a35..b8729933 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -37,6 +37,7 @@ from zipline.data.minute_bars import BcolzMinuteBarReader from zipline.data.data_portal import DataPortal from zipline.data.us_equity_pricing import BcolzDailyBarWriter from zipline.finance.slippage import FixedSlippage +from zipline.finance.restrictions import NoopRestrictions from zipline.protocol import BarData from zipline.testing import ( tmp_trading_env, @@ -317,10 +318,11 @@ class FinanceTestCase(WithLogger, order_date = order_date.replace(hour=14, minute=30) else: bar_data = BarData( - data_portal, - lambda: tick, - sim_params.data_frequency, - self.trading_calendar + data_portal=data_portal, + simulation_dt_func=lambda: tick, + data_frequency=sim_params.data_frequency, + trading_calendar=self.trading_calendar, + restrictions=NoopRestrictions(), ) txns, _, closed_orders = blotter.get_transactions(bar_data) for txn in txns: diff --git a/tests/test_history.py b/tests/test_history.py index a68b7216..2859bb8b 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -21,20 +21,21 @@ import pandas as pd from six import iteritems from zipline import TradingAlgorithm -from zipline._protocol import handle_non_market_minutes +from zipline._protocol import handle_non_market_minutes, BarData from zipline.assets import Asset from zipline.errors import ( HistoryInInitialize, HistoryWindowStartsBeforeData, ) from zipline.finance.trading import SimulationParameters -from zipline.protocol import BarData +from zipline.finance.restrictions import NoopRestrictions from zipline.testing import ( create_minute_df_for_asset, str_to_seconds, MockDailyBarReader, ) from zipline.testing.fixtures import ( + WithCreateBarData, WithDataPortal, ZiplineTestCase, alias, @@ -46,7 +47,7 @@ OHLCP = OHLC + ['price'] ALL_FIELDS = OHLCP + ['volume'] -class WithHistory(WithDataPortal): +class WithHistory(WithCreateBarData, WithDataPortal): TRADING_START_DT = TRADING_ENV_MIN_DATE = START_DATE = pd.Timestamp( '2014-01-03', tz='UTC', @@ -251,8 +252,9 @@ class WithHistory(WithDataPortal): fields = fields if fields is not None else ALL_FIELDS assets = assets if assets is not None else [self.ASSET2, self.ASSET3] - bar_data = BarData(self.data_portal, lambda: dt, mode, - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: dt, + ) check_internal_consistency( bar_data, assets, fields, 10, freq ) @@ -704,8 +706,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): )[0:60] for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, 'minute', - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute, + ) check_internal_consistency( bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, '1m' ) @@ -766,13 +769,12 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): ) )[1] - midnight_bar_data = \ - BarData(self.data_portal, lambda: midnight, 'minute', - self.trading_calendar) - - yesterday_bar_data = \ - BarData(self.data_portal, lambda: last_minute, 'minute', - self.trading_calendar) + midnight_bar_data = self.create_bardata( + lambda: midnight, + ) + yesterday_bar_data = self.create_bardata( + lambda: last_minute + ) with handle_non_market_minutes(midnight_bar_data): for field in ALL_FIELDS: @@ -789,8 +791,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): )[0:60] for idx, minute in enumerate(minutes): - bar_data = BarData(self.data_portal, lambda: minute, 'minute', - self.trading_calendar) + bar_data = self.create_bardata( + lambda: minute + ) check_internal_consistency( bar_data, self.SHORT_ASSET, ALL_FIELDS, 30, '1m' ) @@ -799,8 +802,13 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): data_portal = self.make_data_portal() # choose a window that contains the last minute of the asset - bar_data = BarData(data_portal, lambda: minutes[15], 'minute', - self.trading_calendar) + bar_data = BarData( + data_portal=data_portal, + simulation_dt_func=lambda: minutes[15], + data_frequency='minute', + restrictions=NoopRestrictions(), + trading_calendar=self.trading_calendar, + ) # close high low open price volume # 2015-01-06 20:47:00+00:00 768 770 767 769 768 76800 @@ -1012,8 +1020,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): def test_passing_iterable_to_history_regular_hours(self): # regular hours current_dt = pd.Timestamp("2015-01-06 9:45", tz='US/Eastern') - bar_data = BarData(self.data_portal, lambda: current_dt, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: current_dt, + ) bar_data.history(pd.Index([self.ASSET1, self.ASSET2]), "high", 5, "1m") @@ -1021,8 +1030,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): def test_passing_iterable_to_history_bts(self): # before market hours current_dt = pd.Timestamp("2015-01-07 8:45", tz='US/Eastern') - bar_data = BarData(self.data_portal, lambda: current_dt, "minute", - self.trading_calendar) + bar_data = self.create_bardata( + lambda: current_dt, + ) with handle_non_market_minutes(bar_data): bar_data.history(pd.Index([self.ASSET1, self.ASSET2]), @@ -1031,8 +1041,9 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): def test_overnight_adjustments(self): # Should incorporate adjustments on midnight 01/06 current_dt = pd.Timestamp('2015-01-06 8:45', tz='US/Eastern') - bar_data = BarData(self.data_portal, lambda: current_dt, 'minute', - self.trading_calendar) + bar_data = self.create_bardata( + lambda: current_dt, + ) adj_expected = { 'open': np.arange(8381, 8391) / 4.0, @@ -1341,6 +1352,8 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): + CREATE_BARDATA_DATA_FREQUENCY = 'daily' + @classmethod def make_equity_daily_bar_data(cls): yield 1, cls.create_df_for_asset( @@ -1403,8 +1416,9 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): ) for idx, day in enumerate(days): - bar_data = BarData(self.data_portal, lambda: day, 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: day, + ) check_internal_consistency( bar_data, [self.ASSET2, self.ASSET3], ALL_FIELDS, 10, '1d' ) @@ -1445,10 +1459,9 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): # asset1 ends on 2016-01-30 # asset2 ends on 2015-12-13 - bar_data = BarData(self.data_portal, - lambda: pd.Timestamp('2016-01-06', tz='UTC'), - 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: pd.Timestamp('2016-01-06', tz='UTC'), + ) for field in OHLCP: window = bar_data.history( @@ -1486,8 +1499,9 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): # days has 1/7, 1/8 for idx, day in enumerate(days): - bar_data = BarData(self.data_portal, lambda: day, 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: day, + ) check_internal_consistency( bar_data, self.SHORT_ASSET, ALL_FIELDS, 2, '1d' ) @@ -1639,10 +1653,10 @@ class DailyEquityHistoryTestCase(WithHistory, ZiplineTestCase): # asset1 ends on 2016-01-30 # asset2 ends on 2016-01-04 - bar_data = BarData(self.data_portal, - lambda: pd.Timestamp('2016-01-06 16:00', tz='UTC'), - 'daily', - self.trading_calendar) + bar_data = self.create_bardata( + simulation_dt_func=lambda: + pd.Timestamp('2016-01-06 16:00', tz='UTC'), + ) for field in OHLCP: window = bar_data.history( diff --git a/tests/test_tradesimulation.py b/tests/test_tradesimulation.py index d3bb8b33..c6027004 100644 --- a/tests/test_tradesimulation.py +++ b/tests/test_tradesimulation.py @@ -24,6 +24,7 @@ from zipline import TradingAlgorithm from zipline.gens.sim_engine import BEFORE_TRADING_START_BAR from zipline.finance.performance import PerformanceTracker +from zipline.finance.restrictions import NoopRestrictions from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.sources.benchmark_source import BenchmarkSource from zipline.test_algorithms import NoopAlgorithm @@ -135,6 +136,7 @@ def initialize(context): self.data_portal, BeforeTradingStartsOnlyClock(dt), algo._create_benchmark_source(), + NoopRestrictions(), None ) diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index f0e2aaa5..8290b5a4 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -36,8 +36,10 @@ from ..utils.classproperty import classproperty from ..utils.final import FinalMeta, final from .core import tmp_asset_finder, make_simple_equity_info from zipline.assets import Equity, Future +from zipline.finance.restrictions import NoopRestrictions from zipline.pipeline import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader +from zipline.protocol import BarData from zipline.utils.calendars import ( get_calendar, register_calendar) @@ -1319,3 +1321,17 @@ class WithResponses(object): self.responses = self.enter_instance_context( responses.RequestsMock(), ) + + +class WithCreateBarData(WithDataPortal): + + CREATE_BARDATA_DATA_FREQUENCY = 'minute' + + def create_bardata(self, simulation_dt_func, restrictions=None): + return BarData( + self.data_portal, + simulation_dt_func, + self.CREATE_BARDATA_DATA_FREQUENCY, + self.trading_calendar, + restrictions or NoopRestrictions() + ) From bf8b030417c4b52b0f152cbd5a0c3e52f667f458 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Wed, 28 Sep 2016 15:21:36 -0400 Subject: [PATCH 05/15] MAINT: Deprecate `set_do_not_order_list` In favor of a new method `set_restrictions` which takes a Restrictions object. Calls to `set_do_not_order_list` should raise a deprecation warning and create an equivalent Restrictions object, with which `set_restrictions` will be called. For convenience, create a RestrictionsSet from which the "restrictions" version of a security list can be accessed --- tests/test_algorithm.py | 44 ++++++++++++++++++++++++---------- tests/test_security_list.py | 30 ++++++++++++++++++----- zipline/algorithm.py | 42 +++++++++++++++++++++++++++----- zipline/test_algorithms.py | 6 +++++ zipline/utils/security_list.py | 5 ++++ 5 files changed, 103 insertions(+), 24 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 32d5e8bc..a448f714 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -79,6 +79,7 @@ from zipline.finance.trading import SimulationParameters from zipline.finance.restrictions import ( Restriction, HistoricalRestrictions, + StaticRestrictions, RESTRICTION_STATES, ) from zipline.testing import ( @@ -127,6 +128,7 @@ from zipline.test_algorithms import ( SetMaxOrderCountAlgorithm, SetMaxOrderSizeAlgorithm, SetDoNotOrderListAlgorithm, + SetAssetRestrictionsAlgorithm, SetMaxLeverageAlgorithm, api_algo, api_get_environment_algo, @@ -2793,14 +2795,14 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): env=self.env) self.check_algo_fails(algo, handle_data, 0) - def test_set_do_not_order_list(self): + def test_set_asset_restrictions(self): def handle_data(algo, data): algo.could_trade = data.can_trade(algo.sid(self.sid)) algo.order(algo.sid(self.sid), 100) algo.order_count += 1 - # set the restricted list to be one sid for the entire simulation, + # Set HistoricalRestrictions for one sid for the entire simulation, # and fail. rlm = HistoricalRestrictions([ Restriction( @@ -2808,20 +2810,20 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): self.sim_params.start_session, RESTRICTION_STATES.FROZEN) ]) - algo = SetDoNotOrderListAlgorithm( + algo = SetAssetRestrictionsAlgorithm( sid=self.sid, - restricted_list=rlm, + restrictions=rlm, sim_params=self.sim_params, env=self.env, ) self.check_algo_fails(algo, handle_data, 0) self.assertFalse(algo.could_trade) - # if the restricted list is a static list, then use a shim. - rlm = [self.sid] - algo = SetDoNotOrderListAlgorithm( + # Set StaticRestrictions for one sid and fail. + rlm = StaticRestrictions([self.sid]) + algo = SetAssetRestrictionsAlgorithm( sid=self.sid, - restricted_list=rlm, + restrictions=rlm, sim_params=self.sim_params, env=self.env, ) @@ -2829,9 +2831,9 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): self.assertFalse(algo.could_trade) # just log an error on the violation if we choose not to fail. - algo = SetDoNotOrderListAlgorithm( + algo = SetAssetRestrictionsAlgorithm( sid=self.sid, - restricted_list=rlm, + restrictions=rlm, sim_params=self.sim_params, env=self.env, on_error='log' @@ -2851,14 +2853,32 @@ class TestTradingControls(WithSimParams, WithDataPortal, ZiplineTestCase): self.sim_params.start_session, RESTRICTION_STATES.FROZEN) for sid in [134, 135, 136] ]) + algo = SetAssetRestrictionsAlgorithm( + sid=self.sid, + restrictions=rlm, + sim_params=self.sim_params, + env=self.env, + ) + self.check_algo_succeeds(algo, handle_data) + self.assertTrue(algo.could_trade) + + def test_set_do_not_order_list(self): + + def handle_data(algo, data): + algo.could_trade = data.can_trade(algo.sid(self.sid)) + algo.order(algo.sid(self.sid), 100) + algo.order_count += 1 + + rlm = [self.sid] algo = SetDoNotOrderListAlgorithm( sid=self.sid, restricted_list=rlm, sim_params=self.sim_params, env=self.env, ) - self.check_algo_succeeds(algo, handle_data) - self.assertTrue(algo.could_trade) + + self.check_algo_fails(algo, handle_data, 0) + self.assertFalse(algo.could_trade) def test_set_max_order_size(self): diff --git a/tests/test_security_list.py b/tests/test_security_list.py index 6dd975c9..b381adc3 100644 --- a/tests/test_security_list.py +++ b/tests/test_security_list.py @@ -2,6 +2,7 @@ from datetime import timedelta import pandas as pd from testfixtures import TempDirectory +from nose_parameterized import parameterized from zipline.algorithm import TradingAlgorithm from zipline.errors import TradingControlViolation @@ -29,7 +30,7 @@ LEVERAGED_ETFS = load_from_directory('leveraged_etf_list') class RestrictedAlgoWithCheck(TradingAlgorithm): def initialize(self, symbol): self.rl = SecurityListSet(self.get_datetime, self.asset_finder) - self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.set_asset_restrictions(self.rl.restrict_leveraged_etfs) self.order_count = 0 self.sid = self.symbol(symbol) @@ -43,6 +44,18 @@ class RestrictedAlgoWithCheck(TradingAlgorithm): class RestrictedAlgoWithoutCheck(TradingAlgorithm): + def initialize(self, symbol): + self.rl = SecurityListSet(self.get_datetime, self.asset_finder) + self.set_asset_restrictions(self.rl.restrict_leveraged_etfs) + self.order_count = 0 + self.sid = self.symbol(symbol) + + def handle_data(self, data): + self.order(self.sid, 100) + self.order_count += 1 + + +class RestrictedAlgoWithoutCheckSetDoNotOrderList(TradingAlgorithm): def initialize(self, symbol): self.rl = SecurityListSet(self.get_datetime, self.asset_finder) self.set_do_not_order_list(self.rl.leveraged_etf_list) @@ -57,7 +70,7 @@ class RestrictedAlgoWithoutCheck(TradingAlgorithm): class IterateRLAlgo(TradingAlgorithm): def initialize(self, symbol): self.rl = SecurityListSet(self.get_datetime, self.asset_finder) - self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.set_asset_restrictions(self.rl.restrict_leveraged_etfs) self.order_count = 0 self.sid = self.symbol(symbol) self.found = False @@ -213,10 +226,15 @@ class SecurityListTestCase(WithLogger, WithTradingCalendars, ZiplineTestCase): env=self.env) algo.run(self.data_portal) - def test_algo_with_rl_violation(self): - algo = RestrictedAlgoWithoutCheck(symbol='BZQ', - sim_params=self.sim_params, - env=self.env) + @parameterized.expand([ + ('using_set_do_not_order_list', + RestrictedAlgoWithoutCheckSetDoNotOrderList), + ('using_set_restrictions', RestrictedAlgoWithoutCheck), + ]) + def test_algo_with_rl_violation(self, name, algo_class): + algo = algo_class(symbol='BZQ', + sim_params=self.sim_params, + env=self.env) with self.assertRaises(TradingControlViolation) as ctx: algo.run(self.data_portal) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index d5dc429f..324f247a 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -2173,15 +2173,45 @@ class TradingAlgorithm(object): restricted_list : container[Asset], SecurityList The assets that cannot be ordered. """ + if isinstance(restricted_list, SecurityList): + warnings.warn( + "`set_do_not_order_list(security_lists.leveraged_etf_list)` " + "is deprecated. Use `set_asset_restrictions(" + "security_lists.restrict_leveraged_etfs)` instead.", + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + restrictions = SecurityListRestrictions(restricted_list) + else: + warnings.warn( + "`set_do_not_order_list(container_of_assets)` is deprecated. " + "Create a zipline.finance.restrictions.StaticRestrictions " + "object with a container of assets and use " + "`set_asset_restrictions(StaticRestrictions(" + "container_of_assets))` instead.", + category=ZiplineDeprecationWarning, + stacklevel=2 + ) + restrictions = StaticRestrictions(restricted_list) - if isinstance(restricted_list, (list, tuple, set)): - restricted_list = StaticRestrictions(restricted_list) - elif isinstance(restricted_list, SecurityList): - restricted_list = SecurityListRestrictions(restricted_list) + self.set_asset_restrictions(restrictions, on_error) - control = RestrictedListOrder(on_error, restricted_list) + @api_method + @expect_types( + restrictions=Restrictions, + on_error=str, + ) + def set_asset_restrictions(self, restrictions, on_error='fail'): + """Set a restriction on which assets can be ordered. + + Parameters + ---------- + restricted_list : container[Asset], SecurityList, Restrictions + The assets that cannot be ordered. + """ + control = RestrictedListOrder(on_error, restrictions) self.register_trading_control(control) - self.restrictions = restricted_list + self.restrictions = restrictions @api_method def set_long_only(self, on_error='fail'): diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 052dd821..41121195 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -510,6 +510,12 @@ class SetDoNotOrderListAlgorithm(TradingAlgorithm): self.set_do_not_order_list(restricted_list, on_error) +class SetAssetRestrictionsAlgorithm(TradingAlgorithm): + def initialize(self, sid=None, restrictions=None, on_error='fail'): + self.order_count = 0 + self.set_asset_restrictions(restrictions, on_error) + + class SetMaxOrderCountAlgorithm(TradingAlgorithm): def initialize(self, count): self.order_count = 0 diff --git a/zipline/utils/security_list.py b/zipline/utils/security_list.py index 8532d3fb..e871dedd 100644 --- a/zipline/utils/security_list.py +++ b/zipline/utils/security_list.py @@ -8,6 +8,7 @@ import pytz import zipline from zipline.errors import SymbolNotFound +from zipline.finance.restrictions import SecurityListRestrictions from zipline.zipline_warnings import ZiplineDeprecationWarning @@ -114,6 +115,10 @@ class SecurityListSet(object): ) return self._leveraged_etf + @property + def restrict_leveraged_etfs(self): + return SecurityListRestrictions(self.leveraged_etf_list) + def load_from_directory(list_name): """ From 523b05ef3f6b797796441626897b22cfda054008 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Wed, 28 Sep 2016 18:56:03 -0400 Subject: [PATCH 06/15] TEST: Clarify test_restrictions a bit. - Use parameter_space instead of `parameterized.expand`. - Use a timedelta instead of concatenating strings. - Use a (possibly no-op) scramble function instead of reordering list literals. - Use `freeze_dt, unfreeze_dt, re_freeze_dt` instead of `dts[n]`. - Rename `assert_vectorized_results` to `assert_all_restrictions`. --- tests/test_restrictions.py | 193 ++++++++++++++++++++----------------- 1 file changed, 103 insertions(+), 90 deletions(-) diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py index 74327ce9..078b9699 100644 --- a/tests/test_restrictions.py +++ b/tests/test_restrictions.py @@ -1,9 +1,10 @@ import pandas as pd from pandas.util.testing import assert_series_equal -from nose_parameterized import parameterized from six import iteritems from functools import partial +from toolz import groupby + from zipline.finance.restrictions import ( RESTRICTION_STATES, Restriction, @@ -13,6 +14,7 @@ from zipline.finance.restrictions import ( NoopRestrictions, ) +from zipline.testing import parameter_space from zipline.testing.fixtures import ( WithDataPortal, ZiplineTestCase, @@ -34,6 +36,7 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): 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)) @@ -41,119 +44,122 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): def assert_not_restricted(self, rl, asset, dt): self.assertFalse(rl.is_restricted(asset, dt)) - def assert_vectorized_results(self, rl, expected, 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([self.ASSET1, self.ASSET2, self.ASSET3], dt), - pd.Series( - index=pd.Index([self.ASSET1, self.ASSET2, self.ASSET3]), - data=expected - ) + rl.is_restricted(assets, dt), + pd.Series(index=pd.Index(assets), data=expected), ) - @parameterized.expand([ - ('_'.join([timing, ordering]), - timing == 'intraday', - ordering == 'ordered') - for timing in ['intraday', 'interday'] - for ordering in ['ordered', 'unordered'] - ]) - def test_historical_restrictions(self, name, is_intraday, is_ordered): + @parameter_space( + date_offset=( + pd.Timedelta(0), + pd.Timedelta('1 minute'), + pd.Timedelta('15 hours 5 minutes') + ), + check_unordered=(False, True), + __fail_fast=True, + ) + def test_historical_restrictions(self, date_offset, check_unordered): """ 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 """ - - hour_of_day = ' 15:00' if is_intraday else '' - - if is_ordered: - restriction_dates = { - self.ASSET1: [ - (str_to_ts('2011-01-04' + hour_of_day), FROZEN), - (str_to_ts('2011-01-05' + hour_of_day), ALLOWED), - (str_to_ts('2011-01-06' + hour_of_day), FROZEN), - ], - self.ASSET2: [ - (str_to_ts('2011-01-05' + hour_of_day), FROZEN), - (str_to_ts('2011-01-06' + hour_of_day), ALLOWED), - (str_to_ts('2011-01-07' + hour_of_day), FROZEN), - ], - } + if check_unordered: + def maybe_scramble(rs): + # Swap the first two restrictions to check that we don't care + # that the restriction dates are ordered. + tmp = rs[0] + rs[0] = rs[1] + rs[1] = tmp + return rs else: - restriction_dates = { - self.ASSET1: [ - (str_to_ts('2011-01-05' + hour_of_day), ALLOWED), - (str_to_ts('2011-01-06' + hour_of_day), FROZEN), - (str_to_ts('2011-01-04' + hour_of_day), FROZEN), - ], - self.ASSET2: [ - (str_to_ts('2011-01-06' + hour_of_day), ALLOWED), - (str_to_ts('2011-01-05' + hour_of_day), FROZEN), - (str_to_ts('2011-01-07' + hour_of_day), FROZEN), - ], - } + maybe_scramble = lambda r: r - restrictions = sum([ - [Restriction(asset, info[0], info[1]) for info in r_history] - for asset, r_history in iteritems(restriction_dates) - ], []) - rl = HistoricalRestrictions(restrictions) + 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 + all_restrictions = ( + maybe_scramble([ + Restriction(self.ASSET1, rdate('2011-01-04'), FROZEN), + Restriction(self.ASSET1, rdate('2011-01-05'), ALLOWED), + Restriction(self.ASSET1, rdate('2011-01-06'), FROZEN), + ]) + + + maybe_scramble([ + Restriction(self.ASSET2, rdate('2011-01-05'), FROZEN), + Restriction(self.ASSET2, rdate('2011-01-06'), ALLOWED), + Restriction(self.ASSET2, rdate('2011-01-07'), FROZEN), + ]) + ) + 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_vectorized_results = partial(self.assert_vectorized_results, rl) + assert_all_restrictions = partial(self.assert_all_restrictions, rl) - for asset, r_history in iteritems(restriction_dates): - dts = sorted([info[0] for info in r_history]) + # 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]) + ) - # Not restricted until on or after the freeze - assert_not_restricted(asset, dts[0] - MINUTE) - assert_is_restricted(asset, dts[0]) - assert_is_restricted(asset, dts[0] + MINUTE) + # 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, dts[1] - MINUTE) - assert_not_restricted(asset, dts[1]) - assert_not_restricted(asset, dts[1] + MINUTE) + 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, dts[2] - MINUTE) - assert_is_restricted(asset, dts[2]) - assert_is_restricted(asset, dts[2] + MINUTE) + 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, dts[2] + MINUTE * 1000000) - - dts = [str_to_ts(ts + hour_of_day) for ts in ['2011-01-04', - '2011-01-05', - '2011-01-06', - '2011-01-07']] + 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 - assert_vectorized_results([False, False, False], dts[0] - MINUTE) - assert_vectorized_results([True, False, False], dts[0]) - assert_vectorized_results([True, False, False], dts[0] + MINUTE) + 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 - assert_vectorized_results([True, False, False], dts[1] - MINUTE) - assert_vectorized_results([False, True, False], dts[1]) - assert_vectorized_results([False, True, False], dts[1] + MINUTE) + 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 - assert_vectorized_results([False, True, False], dts[2] - MINUTE) - assert_vectorized_results([True, False, False], dts[2]) - assert_vectorized_results([True, False, False], dts[2] + MINUTE) + 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 - assert_vectorized_results([True, False, False], dts[3] - MINUTE) - assert_vectorized_results([True, True, False], dts[3]) - assert_vectorized_results([True, True, False], dts[3] + MINUTE) + 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_vectorized_results( + assert_all_restrictions( [True, True, False], - dts[3] + MINUTE * 10000000 + d3 + (MINUTE * 10000000) ) def test_historical_restrictions_consecutive_states(self): @@ -202,16 +208,17 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): 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_vectorized_results = partial(self.assert_vectorized_results, 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_vectorized_results([True, True, False], dt) + assert_all_restrictions([True, True, False], dt) def test_security_list_restrictions(self): """ @@ -237,23 +244,29 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): assert_not_restricted = partial(self.assert_not_restricted, rl) assert_is_restricted = partial(self.assert_is_restricted, rl) - assert_vectorized_results = partial(self.assert_vectorized_results, 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_vectorized_results( - [True, False, False], 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_vectorized_results([False, True, True], 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_vectorized_results([True, True, True], str_to_ts('2011-01-05')) + assert_all_restrictions( + [True, True, True], + str_to_ts('2011-01-05') + ) def test_noop_restrictions(self): """ @@ -262,7 +275,7 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): rl = NoopRestrictions() assert_not_restricted = partial(self.assert_not_restricted, rl) - assert_vectorized_results = partial(self.assert_vectorized_results, 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', @@ -270,4 +283,4 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): assert_not_restricted(self.ASSET1, dt) assert_not_restricted(self.ASSET2, dt) assert_not_restricted(self.ASSET3, dt) - assert_vectorized_results([False, False, False], dt) + assert_all_restrictions([False, False, False], dt) From cfce14ed9b6ab4827566f50b99dcebf18d426fc2 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Wed, 28 Sep 2016 19:15:15 -0400 Subject: [PATCH 07/15] DOC: Update docstring for set_restrictions. --- zipline/algorithm.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 324f247a..55f06cf9 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -2206,8 +2206,12 @@ class TradingAlgorithm(object): Parameters ---------- - restricted_list : container[Asset], SecurityList, Restrictions - The assets that cannot be ordered. + restricted_list : Restrictions + An object providing information about restricted assets. + + See Also + -------- + zipline.finance.restrictions.Restrictions """ control = RestrictedListOrder(on_error, restrictions) self.register_trading_control(control) From d47144dfb823b0e0aa00bb563d1c4cd2a41c89b1 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Wed, 28 Sep 2016 19:17:01 -0400 Subject: [PATCH 08/15] DOC: Rename NoopRestrictions to NoRestrictions. --- tests/finance/test_slippage.py | 4 ++-- tests/test_finance.py | 4 ++-- tests/test_history.py | 4 ++-- tests/test_restrictions.py | 4 ++-- tests/test_tradesimulation.py | 4 ++-- zipline/algorithm.py | 4 ++-- zipline/finance/restrictions.py | 2 +- zipline/testing/fixtures.py | 4 ++-- 8 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/finance/test_slippage.py b/tests/finance/test_slippage.py index 17846b94..48d33b95 100644 --- a/tests/finance/test_slippage.py +++ b/tests/finance/test_slippage.py @@ -29,7 +29,7 @@ from zipline.finance.slippage import VolumeShareSlippage from zipline.protocol import DATASOURCE_TYPE, BarData from zipline.finance.blotter import Order -from zipline.finance.restrictions import NoopRestrictions +from zipline.finance.restrictions import NoRestrictions from zipline.data.data_portal import DataPortal from zipline.testing import tmp_bcolz_equity_minute_bar_reader from zipline.testing.fixtures import ( @@ -788,7 +788,7 @@ class OrdersStopTestCase(WithSimParams, lambda: dt, self.sim_params.data_frequency, self.trading_calendar, - NoopRestrictions(), + NoRestrictions(), ) _, txn = next(slippage_model.simulate( diff --git a/tests/test_finance.py b/tests/test_finance.py index b8729933..488d52c7 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -37,7 +37,7 @@ from zipline.data.minute_bars import BcolzMinuteBarReader from zipline.data.data_portal import DataPortal from zipline.data.us_equity_pricing import BcolzDailyBarWriter from zipline.finance.slippage import FixedSlippage -from zipline.finance.restrictions import NoopRestrictions +from zipline.finance.restrictions import NoRestrictions from zipline.protocol import BarData from zipline.testing import ( tmp_trading_env, @@ -322,7 +322,7 @@ class FinanceTestCase(WithLogger, simulation_dt_func=lambda: tick, data_frequency=sim_params.data_frequency, trading_calendar=self.trading_calendar, - restrictions=NoopRestrictions(), + restrictions=NoRestrictions(), ) txns, _, closed_orders = blotter.get_transactions(bar_data) for txn in txns: diff --git a/tests/test_history.py b/tests/test_history.py index 2859bb8b..d1ae15eb 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -28,7 +28,7 @@ from zipline.errors import ( HistoryWindowStartsBeforeData, ) from zipline.finance.trading import SimulationParameters -from zipline.finance.restrictions import NoopRestrictions +from zipline.finance.restrictions import NoRestrictions from zipline.testing import ( create_minute_df_for_asset, str_to_seconds, @@ -806,7 +806,7 @@ class MinuteEquityHistoryTestCase(WithHistory, ZiplineTestCase): data_portal=data_portal, simulation_dt_func=lambda: minutes[15], data_frequency='minute', - restrictions=NoopRestrictions(), + restrictions=NoRestrictions(), trading_calendar=self.trading_calendar, ) diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py index 078b9699..aa6d50cb 100644 --- a/tests/test_restrictions.py +++ b/tests/test_restrictions.py @@ -11,7 +11,7 @@ from zipline.finance.restrictions import ( HistoricalRestrictions, StaticRestrictions, SecurityListRestrictions, - NoopRestrictions, + NoRestrictions, ) from zipline.testing import parameter_space @@ -273,7 +273,7 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): Test single- and multi-asset queries on no-op restrictions """ - rl = NoopRestrictions() + rl = NoRestrictions() assert_not_restricted = partial(self.assert_not_restricted, rl) assert_all_restrictions = partial(self.assert_all_restrictions, rl) diff --git a/tests/test_tradesimulation.py b/tests/test_tradesimulation.py index c6027004..428d409b 100644 --- a/tests/test_tradesimulation.py +++ b/tests/test_tradesimulation.py @@ -24,7 +24,7 @@ from zipline import TradingAlgorithm from zipline.gens.sim_engine import BEFORE_TRADING_START_BAR from zipline.finance.performance import PerformanceTracker -from zipline.finance.restrictions import NoopRestrictions +from zipline.finance.restrictions import NoRestrictions from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.sources.benchmark_source import BenchmarkSource from zipline.test_algorithms import NoopAlgorithm @@ -136,7 +136,7 @@ def initialize(context): self.data_portal, BeforeTradingStartsOnlyClock(dt), algo._create_benchmark_source(), - NoopRestrictions(), + NoRestrictions(), None ) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 55f06cf9..9e2e579c 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -83,7 +83,7 @@ from zipline.finance.slippage import ( ) from zipline.finance.cancel_policy import NeverCancel, CancelPolicy from zipline.finance.restrictions import ( - NoopRestrictions, + NoRestrictions, StaticRestrictions, SecurityListRestrictions, ) @@ -425,7 +425,7 @@ class TradingAlgorithm(object): # A dictionary of the actual capital change deltas, keyed by timestamp self.capital_change_deltas = {} - self.restrictions = NoopRestrictions() + self.restrictions = NoRestrictions() def init_engine(self, get_loader): """ diff --git a/zipline/finance/restrictions.py b/zipline/finance/restrictions.py index 52d3b150..463eb916 100644 --- a/zipline/finance/restrictions.py +++ b/zipline/finance/restrictions.py @@ -48,7 +48,7 @@ class Restrictions(with_metaclass(abc.ABCMeta)): raise NotImplementedError('is_restricted') -class NoopRestrictions(Restrictions): +class NoRestrictions(Restrictions): """ A no-op restrictions that contains no restrictions """ diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index 8290b5a4..eb6bd0f9 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -36,7 +36,7 @@ from ..utils.classproperty import classproperty from ..utils.final import FinalMeta, final from .core import tmp_asset_finder, make_simple_equity_info from zipline.assets import Equity, Future -from zipline.finance.restrictions import NoopRestrictions +from zipline.finance.restrictions import NoRestrictions from zipline.pipeline import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader from zipline.protocol import BarData @@ -1333,5 +1333,5 @@ class WithCreateBarData(WithDataPortal): simulation_dt_func, self.CREATE_BARDATA_DATA_FREQUENCY, self.trading_calendar, - restrictions or NoopRestrictions() + restrictions or NoRestrictions() ) From 87daa75c0ca0c6e00f61273dc5d27e8d25bbe1cf Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Thu, 29 Sep 2016 10:42:23 -0400 Subject: [PATCH 09/15] MAINT: Use Timedelta instead of DateOffset. In days_at_time, use a Timedelta instead of a DateOffset. We were previously using DateOffset to work around a pandas 16 bug even though it raises a PerformanceWarning. Now that we're on pandas 18, we can use the much simpler Timedelta construction. --- zipline/utils/calendars/trading_calendar.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index a8e35a07..a931fd32 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -25,7 +25,6 @@ from pandas import ( DataFrame, date_range, DatetimeIndex, - DateOffset ) from pandas.tseries.offsets import CustomBusinessDay from zipline.utils.calendars._calendar_helpers import ( @@ -830,16 +829,13 @@ def days_at_time(days, t, tz, day_offset=0): # Offset days without tz to avoid timezone issues. days = DatetimeIndex(days).tz_localize(None) - days_offset = days + DateOffset(days=day_offset) - - # Shift all days to the target time in the local timezone, then - # convert to UTC. - - # FIXME: Once we're off Pandas 16, see if we can replace DateOffset with - # TimeDelta. - return days_offset.shift( - 1, freq=DateOffset(hour=t.hour, minute=t.minute, second=t.second) - ).tz_localize(tz).tz_convert('UTC') + delta = pd.Timedelta( + days=day_offset, + hours=t.hour, + minutes=t.minute, + seconds=t.second, + ) + return (days + delta).tz_localize(tz) def holidays_at_time(calendar, start, end, time, tz): From d5ea9c4daa6a3352d7469d988719de4796a927fc Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Thu, 29 Sep 2016 11:31:19 -0400 Subject: [PATCH 10/15] BUG: `days_at_time` should return UTC dates. Adds an example and clarifies the docs. --- zipline/utils/calendars/trading_calendar.py | 27 ++++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index a931fd32..9687103f 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -809,20 +809,35 @@ class TradingCalendar(with_metaclass(ABCMeta)): def days_at_time(days, t, tz, day_offset=0): """ - Shift an index of days to time t, interpreted in tz. + Create an index of days at time ``t``, interpreted in timezone ``tz``. - Overwrites any existing tz info on the input. + The returned index is localized to UTC. Parameters ---------- days : DatetimeIndex - The "base" time which we want to change. + An index of dates (represented as midnight). t : datetime.time - The time we want to offset @days by + The time to apply as an offset to each day in ``days``. tz : pytz.timezone - The timezone which these times represent + The timezone to use to interpret ``tz``. day_offset : int The number of days we want to offset @days by + + Example + ------- + + In the example below, the times switch from 13:45 to 12:45 UTC because + March 13th is the daylight savings transition for US/Eastern. All the + times are still 8:45 when interpreted in US/Eastern. + + >>> import pandas as pd; import datetime; import pprint + >>> dts = pd.date_range('2016-03-12', '2016-03-14') + >>> dts_at_845 = days_at_time(dts, datetime.time(8, 45), 'US/Eastern') + >>> pprint.pprint([str(dt) for dt in dts_at_845]) + ['2016-03-12 13:45:00+00:00', + '2016-03-13 12:45:00+00:00', + '2016-03-14 12:45:00+00:00'] """ if len(days) == 0: return days @@ -835,7 +850,7 @@ def days_at_time(days, t, tz, day_offset=0): minutes=t.minute, seconds=t.second, ) - return (days + delta).tz_localize(tz) + return (days + delta).tz_localize(tz).tz_convert('UTC') def holidays_at_time(calendar, start, end, time, tz): From 5b911756515d1f90a34f76f90f849f5f052dc834 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Thu, 29 Sep 2016 11:33:14 -0400 Subject: [PATCH 11/15] DOC: Fix docstring typo. --- zipline/utils/calendars/trading_calendar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index 9687103f..74239451 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -820,7 +820,7 @@ def days_at_time(days, t, tz, day_offset=0): t : datetime.time The time to apply as an offset to each day in ``days``. tz : pytz.timezone - The timezone to use to interpret ``tz``. + The timezone to use to interpret ``t``. day_offset : int The number of days we want to offset @days by From 7700987abc4bb42df6adfd7e60aeaebcfdcb0dd3 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Thu, 29 Sep 2016 11:33:41 -0400 Subject: [PATCH 12/15] STY: Remove extra whitespace. --- zipline/utils/calendars/trading_calendar.py | 1 - 1 file changed, 1 deletion(-) diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index 74239451..0197ece1 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -826,7 +826,6 @@ def days_at_time(days, t, tz, day_offset=0): Example ------- - In the example below, the times switch from 13:45 to 12:45 UTC because March 13th is the daylight savings transition for US/Eastern. All the times are still 8:45 when interpreted in US/Eastern. From 3b5031a82959371b6c303f2381a7c15acae34e5b Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Thu, 29 Sep 2016 14:54:23 -0400 Subject: [PATCH 13/15] MAINT: Rename restrictions.py to asset_restrictions.py For clarity as to what sort of restrictions these are --- tests/finance/test_slippage.py | 2 +- tests/test_algorithm.py | 2 +- tests/test_bar_data.py | 2 +- tests/test_finance.py | 2 +- tests/test_history.py | 2 +- tests/test_restrictions.py | 2 +- tests/test_tradesimulation.py | 2 +- zipline/_protocol.pyx | 2 +- zipline/algorithm.py | 10 +++++----- zipline/api.py | 4 +--- .../finance/{restrictions.py => asset_restrictions.py} | 0 zipline/finance/controls.py | 2 +- zipline/testing/fixtures.py | 2 +- zipline/utils/security_list.py | 2 +- 14 files changed, 17 insertions(+), 19 deletions(-) rename zipline/finance/{restrictions.py => asset_restrictions.py} (100%) diff --git a/tests/finance/test_slippage.py b/tests/finance/test_slippage.py index 48d33b95..ba1ac919 100644 --- a/tests/finance/test_slippage.py +++ b/tests/finance/test_slippage.py @@ -29,7 +29,7 @@ from zipline.finance.slippage import VolumeShareSlippage from zipline.protocol import DATASOURCE_TYPE, BarData from zipline.finance.blotter import Order -from zipline.finance.restrictions import NoRestrictions +from zipline.finance.asset_restrictions import NoRestrictions from zipline.data.data_portal import DataPortal from zipline.testing import tmp_bcolz_equity_minute_bar_reader from zipline.testing.fixtures import ( diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index a448f714..bfe730e1 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -76,7 +76,7 @@ from zipline.finance.commission import PerShare from zipline.finance.execution import LimitOrder from zipline.finance.order import ORDER_STATUS from zipline.finance.trading import SimulationParameters -from zipline.finance.restrictions import ( +from zipline.finance.asset_restrictions import ( Restriction, HistoricalRestrictions, StaticRestrictions, diff --git a/tests/test_bar_data.py b/tests/test_bar_data.py index b4b1e170..e3ef7edc 100644 --- a/tests/test_bar_data.py +++ b/tests/test_bar_data.py @@ -23,7 +23,7 @@ import pandas as pd from zipline._protocol import handle_non_market_minutes -from zipline.finance.restrictions import ( +from zipline.finance.asset_restrictions import ( Restriction, HistoricalRestrictions, RESTRICTION_STATES, diff --git a/tests/test_finance.py b/tests/test_finance.py index 488d52c7..bd84ba9b 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -37,7 +37,7 @@ from zipline.data.minute_bars import BcolzMinuteBarReader from zipline.data.data_portal import DataPortal from zipline.data.us_equity_pricing import BcolzDailyBarWriter from zipline.finance.slippage import FixedSlippage -from zipline.finance.restrictions import NoRestrictions +from zipline.finance.asset_restrictions import NoRestrictions from zipline.protocol import BarData from zipline.testing import ( tmp_trading_env, diff --git a/tests/test_history.py b/tests/test_history.py index d1ae15eb..233d879d 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -28,7 +28,7 @@ from zipline.errors import ( HistoryWindowStartsBeforeData, ) from zipline.finance.trading import SimulationParameters -from zipline.finance.restrictions import NoRestrictions +from zipline.finance.asset_restrictions import NoRestrictions from zipline.testing import ( create_minute_df_for_asset, str_to_seconds, diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py index aa6d50cb..64e3f500 100644 --- a/tests/test_restrictions.py +++ b/tests/test_restrictions.py @@ -5,7 +5,7 @@ from functools import partial from toolz import groupby -from zipline.finance.restrictions import ( +from zipline.finance.asset_restrictions import ( RESTRICTION_STATES, Restriction, HistoricalRestrictions, diff --git a/tests/test_tradesimulation.py b/tests/test_tradesimulation.py index 428d409b..6277b8e2 100644 --- a/tests/test_tradesimulation.py +++ b/tests/test_tradesimulation.py @@ -24,7 +24,7 @@ from zipline import TradingAlgorithm from zipline.gens.sim_engine import BEFORE_TRADING_START_BAR from zipline.finance.performance import PerformanceTracker -from zipline.finance.restrictions import NoRestrictions +from zipline.finance.asset_restrictions import NoRestrictions from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.sources.benchmark_source import BenchmarkSource from zipline.test_algorithms import NoopAlgorithm diff --git a/zipline/_protocol.pyx b/zipline/_protocol.pyx index b15159fd..43e7d20d 100644 --- a/zipline/_protocol.pyx +++ b/zipline/_protocol.pyx @@ -153,7 +153,7 @@ cdef class BarData: data_frequency : {'minute', 'daily'} The frequency of the bar data; i.e. whether the data is daily or minute bars - restrictions : zipline.finance.restrictions.Restrictions + restrictions : zipline.finance.asset_restrictions.Restrictions Object that combines and returns restricted list information from multiple sources universe_func : callable, optional diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 9e2e579c..df437706 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -76,13 +76,13 @@ from zipline.finance.execution import ( StopOrder, ) from zipline.finance.performance import PerformanceTracker -from zipline.finance.restrictions import Restrictions +from zipline.finance.asset_restrictions import Restrictions from zipline.finance.slippage import ( VolumeShareSlippage, SlippageModel ) from zipline.finance.cancel_policy import NeverCancel, CancelPolicy -from zipline.finance.restrictions import ( +from zipline.finance.asset_restrictions import ( NoRestrictions, StaticRestrictions, SecurityListRestrictions, @@ -2185,8 +2185,8 @@ class TradingAlgorithm(object): else: warnings.warn( "`set_do_not_order_list(container_of_assets)` is deprecated. " - "Create a zipline.finance.restrictions.StaticRestrictions " - "object with a container of assets and use " + "Create a zipline.finance.asset_restrictions." + "StaticRestrictions object with a container of assets and use " "`set_asset_restrictions(StaticRestrictions(" "container_of_assets))` instead.", category=ZiplineDeprecationWarning, @@ -2211,7 +2211,7 @@ class TradingAlgorithm(object): See Also -------- - zipline.finance.restrictions.Restrictions + zipline.finance.asset_restrictions.Restrictions """ control = RestrictedListOrder(on_error, restrictions) self.register_trading_control(control) diff --git a/zipline/api.py b/zipline/api.py index f3e345ac..68b78714 100644 --- a/zipline/api.py +++ b/zipline/api.py @@ -16,12 +16,11 @@ # Note that part of the API is implemented in TradingAlgorithm as # methods (e.g. order). These are added to this namespace via the # decorator ``api_method`` inside of algorithm.py. -from .finance.restrictions import ( +from .finance.asset_restrictions import ( Restriction, StaticRestrictions, HistoricalRestrictions, RESTRICTION_STATES, - SecurityListRestrictions, ) from .finance import commission, execution, slippage, cancel_policy from .finance.cancel_policy import ( @@ -47,7 +46,6 @@ __all__ = [ 'StaticRestrictions', 'HistoricalRestrictions', 'RESTRICTION_STATES', - 'SecurityListRestrictions', 'cancel_policy', 'commission', 'date_rules', diff --git a/zipline/finance/restrictions.py b/zipline/finance/asset_restrictions.py similarity index 100% rename from zipline/finance/restrictions.py rename to zipline/finance/asset_restrictions.py diff --git a/zipline/finance/controls.py b/zipline/finance/controls.py index 95d61c89..ca31a462 100644 --- a/zipline/finance/controls.py +++ b/zipline/finance/controls.py @@ -137,7 +137,7 @@ class RestrictedListOrder(TradingControl): Parameters ---------- - restrictions : zipline.finance.restrictions.Restrictions + restrictions : zipline.finance.asset_restrictions.Restrictions Object representing restrictions of a group of assets. """ diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index eb6bd0f9..3570d255 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -36,7 +36,7 @@ from ..utils.classproperty import classproperty from ..utils.final import FinalMeta, final from .core import tmp_asset_finder, make_simple_equity_info from zipline.assets import Equity, Future -from zipline.finance.restrictions import NoRestrictions +from zipline.finance.asset_restrictions import NoRestrictions from zipline.pipeline import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader from zipline.protocol import BarData diff --git a/zipline/utils/security_list.py b/zipline/utils/security_list.py index e871dedd..2a72952c 100644 --- a/zipline/utils/security_list.py +++ b/zipline/utils/security_list.py @@ -8,7 +8,7 @@ import pytz import zipline from zipline.errors import SymbolNotFound -from zipline.finance.restrictions import SecurityListRestrictions +from zipline.finance.asset_restrictions import SecurityListRestrictions from zipline.zipline_warnings import ZiplineDeprecationWarning From 2104a35af8d43a398b4df4a947d79b642189c346 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Fri, 30 Sep 2016 12:26:59 -0400 Subject: [PATCH 14/15] 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 From c5ee71afe63265082b24e595ab74b0f39c4e5b10 Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Fri, 30 Sep 2016 15:57:20 -0400 Subject: [PATCH 15/15] DOC: Clean up Restrictions documentation --- zipline/finance/asset_restrictions.py | 49 +++++++++++++++------------ 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/zipline/finance/asset_restrictions.py b/zipline/finance/asset_restrictions.py index 20b7eb73..c22d5086 100644 --- a/zipline/finance/asset_restrictions.py +++ b/zipline/finance/asset_restrictions.py @@ -25,7 +25,8 @@ RESTRICTION_STATES = enum( class Restrictions(with_metaclass(abc.ABCMeta)): """ - Abstract restricted list interface + Abstract restricted list interface, representing a set of assets that an + algorithm is restricted from trading. """ @abc.abstractmethod @@ -49,31 +50,34 @@ class Restrictions(with_metaclass(abc.ABCMeta)): raise NotImplementedError('is_restricted') def __or__(self, other_restriction): + """Base implementation for combining two restrictions. """ - Base implementation for combining two restrictions. If the right side - is a _UnionRestrictions, calls the overriding implementation with - _UnionRestrictions on the left side - """ + # If the right side is a _UnionRestrictions, defers to the + # _UnionRestrictions implementation of `|`, which intelligently + # flattens restricted lists if isinstance(other_restriction, _UnionRestrictions): - return _UnionRestrictions.__or__(other_restriction, self) + return other_restriction | self return _UnionRestrictions([self, other_restriction]) class _UnionRestrictions(Restrictions): """ - A union of a number of sub restrictions + A union of a number of sub restrictions. Parameters ---------- sub_restrictions : iterable of Restrictions (but not _UnionRestrictions) The Restrictions to be added together + + Notes + ----- + - Consumers should not construct instances of this class directly, but + instead use the `|` operator to combine restrictions """ def __new__(cls, sub_restrictions): - """ - Returns a _UnionRestrictions defined by a list of sub_restrictions, - while dealing with trivial NoRestrictions cases - """ + # Filter out NoRestrictions and deal with resulting cases involving + # one or zero sub_restrictions sub_restrictions = [ r for r in sub_restrictions if not isinstance(r, NoRestrictions) ] @@ -88,10 +92,10 @@ class _UnionRestrictions(Restrictions): def __or__(self, other_restriction): """ - Overrides the base implementation if the left side is a - _UnionRestrictions. Extracts the underlying sub_restrictions from the - _UnionRestrictions + Overrides the base implementation for combining two restrictions, of + which the left side is a _UnionRestrictions. """ + # Flatten the underlying sub restrictions of _UnionRestrictions if isinstance(other_restriction, _UnionRestrictions): new_sub_restrictions = \ self.sub_restrictions + other_restriction.sub_restrictions @@ -103,7 +107,8 @@ class _UnionRestrictions(Restrictions): def is_restricted(self, assets, dt): if isinstance(assets, Asset): return any( - r.is_restricted(assets, dt) for r in self.sub_restrictions) + r.is_restricted(assets, dt) for r in self.sub_restrictions + ) return reduce( operator.or_, @@ -113,18 +118,18 @@ class _UnionRestrictions(Restrictions): class NoRestrictions(Restrictions): """ - A no-op restrictions that contains no restrictions + A no-op restrictions that contains no restrictions. """ def is_restricted(self, assets, dt): if isinstance(assets, Asset): return False - return pd.Series(index=pd.Index(assets), data=[False]*len(assets)) + return pd.Series(index=pd.Index(assets), data=False) class StaticRestrictions(Restrictions): """ Static restrictions stored in memory that are constant regardless of dt - for each asset + for each asset. Parameters ---------- @@ -137,7 +142,7 @@ class StaticRestrictions(Restrictions): def is_restricted(self, assets, dt): """ - An asset is restricted for all dts if it is in the static list + An asset is restricted for all dts if it is in the static list. """ if isinstance(assets, Asset): return assets in self._restricted_set @@ -150,7 +155,7 @@ class StaticRestrictions(Restrictions): class HistoricalRestrictions(Restrictions): """ Historical restrictions stored in memory with effective dates for each - asset + asset. Parameters ---------- @@ -172,7 +177,7 @@ class HistoricalRestrictions(Restrictions): def is_restricted(self, assets, dt): """ Returns whether or not an asset or iterable of assets is restricted - on a dt + on a dt. """ if isinstance(assets, Asset): return self._is_restricted_for_asset(assets, dt) @@ -194,7 +199,7 @@ class HistoricalRestrictions(Restrictions): class SecurityListRestrictions(Restrictions): """ - Restrictions based on a security list + Restrictions based on a security list. Parameters ----------