diff --git a/MANIFEST.in b/MANIFEST.in index 04f196ac..f93ec43c 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,4 @@ include README.md include LICENSE + +recursive-include zipline/resources *.* \ No newline at end of file diff --git a/setup.py b/setup.py index 1ca18e20..464a2e7a 100644 --- a/setup.py +++ b/setup.py @@ -48,6 +48,7 @@ setup( author_email='opensource@quantopian.com', packages=find_packages(), scripts=['scripts/run_algo.py'], + include_package_data=True, long_description=LONG_DESCRIPTION, license='Apache 2.0', classifiers=[ diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index b9c874e7..ab275647 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -55,6 +55,7 @@ from zipline.test_algorithms import ( SetMaxPositionSizeAlgorithm, SetMaxOrderCountAlgorithm, SetMaxOrderSizeAlgorithm, + SetDoNotOrderListAlgorithm, api_algo, api_get_environment_algo, api_symbol_algo, @@ -955,6 +956,29 @@ class TestTradingControls(TestCase): algo = SetMaxPositionSizeAlgorithm(max_shares=10, max_notional=61.0) 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]) + + def handle_data(algo, data): + algo.order(self.sid, 100) + algo.order_count += 1 + + self.check_algo_fails(algo, handle_data, 0) + + # set the restricted list to exclude the sid, and succeed + algo = SetDoNotOrderListAlgorithm( + sid=self.sid, + restricted_list=[134, 135, 136]) + + def handle_data(algo, data): + algo.order(self.sid, 100) + algo.order_count += 1 + + self.check_algo_succeeds(algo, handle_data) + def test_set_max_order_size(self): # Buy one share. diff --git a/tests/test_security_list.py b/tests/test_security_list.py new file mode 100644 index 00000000..bb54868f --- /dev/null +++ b/tests/test_security_list.py @@ -0,0 +1,318 @@ +import pytz +from datetime import datetime, timedelta +from unittest import TestCase + + +from zipline.algorithm import TradingAlgorithm +from zipline.errors import TradingControlViolation +from zipline.sources import SpecificEquityTrades +from zipline.utils.test_utils import ( + setup_logger, add_security_data, remove_security_data_directory) +from zipline.utils import factory +from zipline.utils.security_list import ( + SecurityListSet, load_from_directory) + +LEVERAGED_ETFS = load_from_directory('leveraged_etf_list') + + +class RestrictedAlgoWithCheck(TradingAlgorithm): + def initialize(self, sid): + self.rl = SecurityListSet(self.get_datetime) + self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.order_count = 0 + self.sid = sid + + def handle_data(self, data): + if not self.order_count: + if self.sid not in \ + self.rl.leveraged_etf_list: + self.order(self.sid, 100) + self.order_count += 1 + + +class RestrictedAlgoWithoutCheck(TradingAlgorithm): + def initialize(self, sid): + self.rl = SecurityListSet(self.get_datetime) + self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.order_count = 0 + self.sid = sid + + def handle_data(self, data): + self.order(self.sid, 100) + self.order_count += 1 + + +class IterateRLAlgo(TradingAlgorithm): + def initialize(self, sid): + self.rl = SecurityListSet(self.get_datetime) + self.set_do_not_order_list(self.rl.leveraged_etf_list) + self.order_count = 0 + self.sid = sid + self.found = False + + def handle_data(self, data): + for stock in self.rl.leveraged_etf_list: + if stock == self.sid: + self.found = True + + +class SecurityListTestCase(TestCase): + + def setUp(self): + self.extra_knowledge_date = \ + datetime(2015, 1, 27, 0, 0, tzinfo=pytz.utc) + self.trading_day_before_first_kd = datetime( + 2015, 1, 23, 0, 0, tzinfo=pytz.utc) + + setup_logger(self) + + def test_iterate_over_rl(self): + sim_params = factory.create_simulation_parameters( + start=list(LEVERAGED_ETFS.keys())[0], num_days=4) + + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + algo = IterateRLAlgo(sid='BZQ', sim_params=sim_params) + algo.run(self.source) + self.assertTrue(algo.found) + + def test_security_list(self): + + # set the knowledge date to the first day of the + # leveraged etf knowledge date. + def get_datetime(): + return list(LEVERAGED_ETFS.keys())[0] + + rl = SecurityListSet(get_datetime) + # assert that a sample from the leveraged list are in restricted + + self.assertIn("BZQ", rl.leveraged_etf_list) + self.assertIn("URTY", rl.leveraged_etf_list) + self.assertIn("JFT", rl.leveraged_etf_list) + + # assert that a sample of allowed stocks are not in restricted + # AAPL + self.assertNotIn("AAPL", rl.leveraged_etf_list) + # GOOG + self.assertNotIn("GOOG", rl.leveraged_etf_list) + + def test_security_add(self): + def get_datetime(): + return datetime(2015, 1, 27, tzinfo=pytz.utc) + try: + add_security_data(['AAPL', 'GOOG'], []) + rl = SecurityListSet(get_datetime) + self.assertIn("AAPL", rl.leveraged_etf_list) + self.assertIn("GOOG", rl.leveraged_etf_list) + self.assertIn("BZQ", rl.leveraged_etf_list) + self.assertIn("URTY", rl.leveraged_etf_list) + finally: + remove_security_data_directory() + + def test_security_add_delete(self): + try: + def get_datetime(): + return datetime(2015, 1, 27, tzinfo=pytz.utc) + add_security_data([], ['BZQ', 'URTY']) + rl = SecurityListSet(get_datetime) + self.assertNotIn("BZQ", rl.leveraged_etf_list) + self.assertNotIn("URTY", rl.leveraged_etf_list) + finally: + remove_security_data_directory() + + def test_algo_without_rl_violation_via_check(self): + sim_params = factory.create_simulation_parameters( + start=list(LEVERAGED_ETFS.keys())[0], num_days=4) + + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + + algo = RestrictedAlgoWithCheck(sid='BZQ', sim_params=sim_params) + algo.run(self.source) + + def test_algo_without_rl_violation(self): + sim_params = factory.create_simulation_parameters( + start=list(LEVERAGED_ETFS.keys())[0], num_days=4) + + trade_history = factory.create_trade_history( + 'AAPL', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + algo = RestrictedAlgoWithoutCheck(sid='AAPL', sim_params=sim_params) + algo.run(self.source) + + def test_algo_with_rl_violation(self): + sim_params = factory.create_simulation_parameters( + start=list(LEVERAGED_ETFS.keys())[0], num_days=4) + + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + + self.df_source, self.df = \ + factory.create_test_df_source(sim_params) + + algo = RestrictedAlgoWithoutCheck(sid='BZQ', sim_params=sim_params) + with self.assertRaises(TradingControlViolation) as ctx: + algo.run(self.source) + + self.check_algo_exception(algo, ctx, 0) + + # repeat with a symbol from a different lookup date + + trade_history = factory.create_trade_history( + 'JFT', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + + self.df_source, self.df = \ + factory.create_test_df_source(sim_params) + + algo = RestrictedAlgoWithoutCheck(sid='JFT', sim_params=sim_params) + with self.assertRaises(TradingControlViolation) as ctx: + algo.run(self.source) + + self.check_algo_exception(algo, ctx, 0) + + def test_algo_with_rl_violation_on_knowledge_date(self): + + sim_params = factory.create_simulation_parameters( + start=self.trading_day_before_first_kd, num_days=4) + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + + algo = RestrictedAlgoWithoutCheck(sid='BZQ', sim_params=sim_params) + with self.assertRaises(TradingControlViolation) as ctx: + algo.run(self.source) + + self.check_algo_exception(algo, ctx, 1) + + def test_algo_with_rl_violation_after_knowledge_date(self): + sim_params = factory.create_simulation_parameters( + start=list( + LEVERAGED_ETFS.keys())[0] + timedelta(days=7), num_days=5) + + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + algo = RestrictedAlgoWithoutCheck(sid='BZQ', sim_params=sim_params) + with self.assertRaises(TradingControlViolation) as ctx: + algo.run(self.source) + + self.check_algo_exception(algo, ctx, 0) + + def test_algo_with_rl_violation_cumulative(self): + """ + Add a new restriction, run a test long after both + knowledge dates, make sure stock from original restriction + set is still disallowed. + """ + sim_params = factory.create_simulation_parameters( + start=list( + LEVERAGED_ETFS.keys())[0] + timedelta(days=7), num_days=4) + + try: + add_security_data(['AAPL'], []) + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + algo = RestrictedAlgoWithoutCheck( + sid='BZQ', sim_params=sim_params) + with self.assertRaises(TradingControlViolation) as ctx: + algo.run(self.source) + + self.check_algo_exception(algo, ctx, 0) + finally: + remove_security_data_directory() + + def test_algo_without_rl_violation_after_delete(self): + try: + # add a delete statement removing bzq + # write a new delete statement file to disk + add_security_data([], ['BZQ']) + + sim_params = factory.create_simulation_parameters( + start=self.extra_knowledge_date, num_days=3) + trade_history = factory.create_trade_history( + 'BZQ', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + algo = RestrictedAlgoWithoutCheck( + sid='BZQ', sim_params=sim_params) + algo.run(self.source) + finally: + remove_security_data_directory() + + def test_algo_with_rl_violation_after_add(self): + try: + add_security_data(['AAPL'], []) + sim_params = factory.create_simulation_parameters( + start=self.trading_day_before_first_kd, num_days=4) + trade_history = factory.create_trade_history( + 'AAPL', + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + sim_params + ) + self.source = SpecificEquityTrades(event_list=trade_history) + algo = RestrictedAlgoWithoutCheck( + sid='AAPL', sim_params=sim_params) + with self.assertRaises(TradingControlViolation) as ctx: + algo.run(self.source) + + self.check_algo_exception(algo, ctx, 2) + finally: + remove_security_data_directory() + + def check_algo_exception(self, algo, ctx, expected_order_count): + self.assertEqual(algo.order_count, expected_order_count) + exc = ctx.exception + self.assertEqual(TradingControlViolation, type(exc)) + exc_msg = str(ctx.exception) + self.assertTrue("RestrictedListOrder" in exc_msg) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 64e9b9a0..9a705d86 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -48,6 +48,7 @@ from zipline.finance.controls import ( MaxOrderCount, MaxOrderSize, MaxPositionSize, + RestrictedListOrder ) from zipline.finance.execution import ( LimitOrder, @@ -1080,6 +1081,14 @@ class TradingAlgorithm(object): control = MaxOrderCount(max_count) self.register_trading_control(control) + @api_method + def set_do_not_order_list(self, restricted_list): + """ + Set a restriction on which sids can be ordered. + """ + control = RestrictedListOrder(restricted_list) + self.register_trading_control(control) + @api_method def set_long_only(self): """ diff --git a/zipline/finance/controls.py b/zipline/finance/controls.py index 4e5f6511..d38953e8 100644 --- a/zipline/finance/controls.py +++ b/zipline/finance/controls.py @@ -99,6 +99,35 @@ class MaxOrderCount(TradingControl): self.orders_placed += 1 +class RestrictedListOrder(TradingControl): + """ + TradingControl representing a restricted list of securities that + cannot be ordered by the algorithm. + """ + + def __init__(self, restricted_list): + """ + restricted list can be an iterable or a + container (implements __contains__) for dynamic + restrictions. + """ + + super(RestrictedListOrder, self).__init__() + self.restricted_list = restricted_list + + def validate(self, + sid, + amount, + _portfolio, + _algo_datetime, + _algo_current_data): + """ + Fail if the sid is in the restricted_list. + """ + if sid in self.restricted_list: + self.fail(sid, amount) + + class MaxOrderSize(TradingControl): """ TradingControl representing a limit on the magnitude of any single order diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20120913/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20120913/add new file mode 100644 index 00000000..2680b739 --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20120913/add @@ -0,0 +1,2 @@ +ROSA +SFSA diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20120913/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20120913/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20120919/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20120919/add new file mode 100644 index 00000000..87935918 --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20120919/add @@ -0,0 +1 @@ +RTSA diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20120919/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20120919/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20121012/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20121012/add new file mode 100644 index 00000000..e2f8ae1a --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20121012/add @@ -0,0 +1 @@ +VZZB diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20121012/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20121012/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20130605/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20130605/add new file mode 100644 index 00000000..aefbe5eb --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20130605/add @@ -0,0 +1,2 @@ +LSKY +SSDL diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20130605/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20130605/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20130916/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20130916/add new file mode 100644 index 00000000..e446ea27 --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20130916/add @@ -0,0 +1 @@ +MFSA diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20130916/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20130916/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131002/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131002/add new file mode 100644 index 00000000..7a55ca5a --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131002/add @@ -0,0 +1,2 @@ +DSTJ +DSXJ diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131002/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131002/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131009/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131009/add new file mode 100644 index 00000000..7c5e71b0 --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131009/add @@ -0,0 +1 @@ +JFT diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131009/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131009/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131121/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131121/add new file mode 100644 index 00000000..0d50b50c --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131121/add @@ -0,0 +1,4 @@ +FSA +FSE +FSG +FSU diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131121/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131121/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131227/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131227/add new file mode 100644 index 00000000..e14acb0c --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131227/add @@ -0,0 +1,2 @@ +IPLT +LPLT diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20131227/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20131227/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20140410/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20140410/add new file mode 100644 index 00000000..b403214d --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20140410/add @@ -0,0 +1 @@ +BXDB diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20140410/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20140410/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20140923/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20140923/add new file mode 100644 index 00000000..ba1b836d --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20140923/add @@ -0,0 +1 @@ +BRZS diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20140923/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20140923/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20141119/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20141119/add new file mode 100644 index 00000000..eff691f2 --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20141119/add @@ -0,0 +1,2 @@ +BXUB +BXUC diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20141119/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20141119/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20141226/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20141226/add new file mode 100644 index 00000000..e397a94f --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20141226/add @@ -0,0 +1 @@ +BARS diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20141226/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20141226/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20150123/add b/zipline/resources/security_lists/leveraged_etf_list/20150126/20150123/add new file mode 100644 index 00000000..35f60e65 --- /dev/null +++ b/zipline/resources/security_lists/leveraged_etf_list/20150126/20150123/add @@ -0,0 +1,207 @@ +AGA +AGQ +BAR +BDCL +BDD +BIB +BIS +BOIL +BOM +BRZU +BUNT +BZQ +CEFL +CMD +CROC +CSMB +CURE +DAG +DDM +DEE +DGAZ +DGLD +DGP +DGZ +DIG +DPK +DRN +DRR +DRV +DSLV +DTO +DUG +DUST +DVHL +DVYL +DWTI +DXD +DYY +DZK +DZZ +EDC +EDZ +EET +EEV +EFO +EFU +EMLB +EMSA +EPV +ERX +ERY +EUO +EURL +EWV +EZJ +FAS +FAZ +FBG +FBGX +FEEU +FIBG +FIEG +FIEU +FIGY +FINU +FINZ +FLGE +FXP +GASL +GDAY +GLL +HDLV +IGU +INDL +ITLT +JDST +JGBD +JGBT +JNUG +JPNL +JPX +KOLD +KORU +KRU +LBJ +LBND +LLDM +LLSC +LLSP +LMLP +LTL +MATL +MDLL +MFLA +MIDU +MIDZ +MLPL +MORL +MVV +MZZ +NUGT +PST +QID +QLD +RETL +REW +RGRA +RGRC +RGRE +RGRI +RGRP +ROLA +ROM +RTLA +RUSL +RUSS +RWXL +RXD +RXL +SAA +SBND +SCC +SCO +SDD +SDOW +SDP +SDS +SDYL +SFLA +SIJ +SKF +SMDD +SMK +SMLL +SMN +SOXL +SOXS +SPLX +SPUU +SPXL +SPXS +SPXU +SQQQ +SRS +SRTY +SSG +SSO +SYTL +SZK +TBAR +TBT +TBZ +TECL +TECS +TLL +TMF +TMV +TNA +TPS +TQQQ +TTT +TVIX +TVIZ +TWM +TYD +TYO +TZA +UBR +UBT +UCC +UCD +UCI +UCO +UDNT +UDOW +UGAZ +UGE +UGL +UGLD +UJB +ULE +UMDD +UMX +UPRO +UPV +UPW +URE +URR +URTY +USD +USLV +UST +USV +UUPT +UVXY +UWM +UWTI +UXI +UXJ +UYG +UYM +XPP +YANG +YCL +YCS +YINN +ZSL diff --git a/zipline/resources/security_lists/leveraged_etf_list/20150126/20150123/delete b/zipline/resources/security_lists/leveraged_etf_list/20150126/20150123/delete new file mode 100644 index 00000000..e69de29b diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index d25adbc3..132b6bb9 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -430,6 +430,12 @@ class SetMaxOrderSizeAlgorithm(TradingAlgorithm): max_notional=max_notional) +class SetDoNotOrderListAlgorithm(TradingAlgorithm): + def initialize(self, sid=None, restricted_list=None): + self.order_count = 0 + self.set_do_not_order_list(restricted_list) + + 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 new file mode 100644 index 00000000..4de9dbaf --- /dev/null +++ b/zipline/utils/security_list.py @@ -0,0 +1,141 @@ +import os.path +import pytz +import pandas as pd +from datetime import datetime +from os import listdir + +DATE_FORMAT = "%Y%m%d" +import zipline +zipline_dir = os.path.join(*zipline.__path__) +SECURITY_LISTS_DIR = os.path.join(zipline_dir, 'resources', 'security_lists') + + +def loopback(symbol, *args, **kwargs): + return symbol + + +class SecurityListSet(object): + + def __init__(self, current_date_func, lookup_func=None): + if lookup_func is None: + self.lookup_func = loopback + else: + self.lookup_func = lookup_func + self.current_date_func = current_date_func + self._leveraged_etf = None + + @property + def leveraged_etf_list(self): + if self._leveraged_etf is None: + self._leveraged_etf = SecurityList( + self.lookup_func, + load_from_directory('leveraged_etf_list'), + self.current_date_func + ) + return self._leveraged_etf + + +class SecurityList(object): + + def __init__(self, lookup_func, data, current_date_func): + """ + lookup_func: function that takes a string symbol and a date and + returns a Security object. + data: a nested dictionary: + knowledge_date -> lookup_date -> + {add: [symbol list], 'delete': []}, delete: [symbol list]} + current_date_func: function taking no parameters, returning + current datetime + """ + self.lookup_func = lookup_func + self.data = data + self._cache = {} + self._knowledge_dates = self.make_knowledge_dates(self.data) + self.current_date = current_date_func + self.count = 0 + self._current_set = set() + + def make_knowledge_dates(self, data): + knowledge_dates = sorted( + [pd.Timestamp(k) for k in data.keys()]) + return knowledge_dates + + def __iter__(self): + return iter(self.restricted_list) + + def __contains__(self, item): + return item in self.restricted_list + + @property + def restricted_list(self): + + cd = self.current_date() + for kd in self._knowledge_dates: + if cd < kd: + break + if kd in self._cache: + self._current_set = self._cache[kd] + continue + + for effective_date, changes in iter(self.data[kd].items()): + self.update_current( + effective_date, + changes['add'], + self._current_set.add + ) + + self.update_current( + effective_date, + changes['delete'], + self._current_set.remove + ) + + self._cache[kd] = self._current_set + return self._current_set + + def update_current(self, effective_date, symbols, change_func): + for symbol in symbols: + sid = self.lookup_func( + symbol, + as_of_date=effective_date + ) + change_func(sid) + + +def load_from_directory(list_name): + """ + To resolve the symbol in the LEVERAGED_ETF list, + the date on which the symbol was in effect is needed. + + Furthermore, to maintain a point in time record of our own maintenance + of the restricted list, we need a knowledge date. Thus, restricted lists + are dictionaries of datetime->symbol lists. + new symbols should be entered as a new knowledge date entry. + + This method assumes a directory structure of: + SECURITY_LISTS_DIR/listname/knowledge_date/lookup_date/add.txt + SECURITY_LISTS_DIR/listname/knowledge_date/lookup_date/delete.txt + + The return value is a dictionary with: + knowledge_date -> lookup_date -> + {add: [symbol list], 'delete': [symbol list]} + """ + data = {} + dir_path = os.path.join(SECURITY_LISTS_DIR, list_name) + for kd_name in listdir(dir_path): + kd = datetime.strptime(kd_name, DATE_FORMAT).replace( + tzinfo=pytz.utc) + data[kd] = {} + kd_path = os.path.join(dir_path, kd_name) + for ld_name in listdir(kd_path): + ld = datetime.strptime(ld_name, DATE_FORMAT).replace( + tzinfo=pytz.utc) + data[kd][ld] = {} + ld_path = os.path.join(kd_path, ld_name) + for fname in listdir(ld_path): + fpath = os.path.join(ld_path, fname) + with open(fpath) as f: + symbols = f.read().splitlines() + data[kd][ld][fname] = symbols + + return data diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py index 451590f7..e29f5f7c 100644 --- a/zipline/utils/test_utils.py +++ b/zipline/utils/test_utils.py @@ -1,10 +1,13 @@ from contextlib import contextmanager from logbook import FileHandler from zipline.finance.blotter import ORDER_STATUS +from zipline.utils.security_list import SECURITY_LISTS_DIR from six import itervalues +import os import pandas as pd +import shutil def to_utc(time_str): @@ -129,3 +132,30 @@ def nullctx(): do_stuff() """ yield + + +def add_security_data(adds, deletes): + directory = os.path.join( + SECURITY_LISTS_DIR, + "leveraged_etf_list/20150127/20150125" + ) + if not os.path.exists(directory): + os.makedirs(directory) + del_path = os.path.join(directory, "delete") + with open(del_path, 'w') as f: + for sym in deletes: + f.write(sym) + f.write('\n') + add_path = os.path.join(directory, "add") + with open(add_path, 'w') as f: + for sym in adds: + f.write(sym) + f.write('\n') + + +def remove_security_data_directory(): + directory = os.path.join( + SECURITY_LISTS_DIR, + "leveraged_etf_list/20150127/" + ) + shutil.rmtree(directory)