diff --git a/tests/test_security_list.py b/tests/test_security_list.py index 667b684f..2a18216f 100644 --- a/tests/test_security_list.py +++ b/tests/test_security_list.py @@ -1,6 +1,4 @@ import pytz -import os.path -import shutil from datetime import datetime, timedelta from unittest import TestCase @@ -8,10 +6,11 @@ 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 +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, SECURITY_LISTS_DIR) + SecurityListSet, load_from_directory) LEVERAGED_ETFS = load_from_directory('leveraged_etf_list') @@ -19,14 +18,14 @@ 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.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.rl.leveraged_etf_list: self.order(self.sid, 100) self.order_count += 1 @@ -34,7 +33,7 @@ class RestrictedAlgoWithCheck(TradingAlgorithm): 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.set_do_not_order_list(self.rl.leveraged_etf_list) self.order_count = 0 self.sid = sid @@ -46,13 +45,13 @@ class RestrictedAlgoWithoutCheck(TradingAlgorithm): 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.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: + for stock in self.rl.leveraged_etf_list: if stock == self.sid: self.found = True @@ -93,38 +92,38 @@ class SecurityListTestCase(TestCase): 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("BZQ", rl.leveraged_etf_list) + self.assertIn("URTY", rl.leveraged_etf_list) # assert that a sample of allowed stocks are not in restricted # AAPL - self.assertNotIn("AAPL", rl.LEVERAGED_ETF_LIST) + self.assertNotIn("AAPL", rl.leveraged_etf_list) # GOOG - self.assertNotIn("GOOG", rl.LEVERAGED_ETF_LIST) + 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_data(['AAPL', 'GOOG'], []) + 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) + 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_data_directory() + remove_security_data_directory() def test_security_add_delete(self): try: def get_datetime(): return datetime(2015, 1, 27, tzinfo=pytz.utc) - add_data([], ['BZQ', 'URTY']) + add_security_data([], ['BZQ', 'URTY']) rl = SecurityListSet(get_datetime) - self.assertNotIn("BZQ", rl.LEVERAGED_ETF_LIST) - self.assertNotIn("URTY", rl.LEVERAGED_ETF_LIST) + self.assertNotIn("BZQ", rl.leveraged_etf_list) + self.assertNotIn("URTY", rl.leveraged_etf_list) finally: - remove_data_directory() + remove_security_data_directory() def test_algo_without_rl_violation_via_check(self): sim_params = factory.create_simulation_parameters( @@ -226,7 +225,7 @@ class SecurityListTestCase(TestCase): start=LEVERAGED_ETFS.keys()[0] + timedelta(days=7), num_days=4) try: - add_data(['AAPL'], []) + add_security_data(['AAPL'], []) trade_history = factory.create_trade_history( 'BZQ', [10.0, 10.0, 11.0, 11.0], @@ -242,13 +241,13 @@ class SecurityListTestCase(TestCase): self.check_algo_exception(algo, ctx, 0) finally: - remove_data_directory() + 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_data([], ['BZQ']) + add_security_data([], ['BZQ']) sim_params = factory.create_simulation_parameters( start=self.extra_knowledge_date, num_days=3) @@ -264,11 +263,11 @@ class SecurityListTestCase(TestCase): sid='BZQ', sim_params=sim_params) algo.run(self.source) finally: - remove_data_directory() + remove_security_data_directory() def test_algo_with_rl_violation_after_add(self): try: - add_data(['AAPL'], []) + 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( @@ -286,7 +285,7 @@ class SecurityListTestCase(TestCase): self.check_algo_exception(algo, ctx, 2) finally: - remove_data_directory() + remove_security_data_directory() def check_algo_exception(self, algo, ctx, expected_order_count): self.assertEqual(algo.order_count, expected_order_count) @@ -294,30 +293,3 @@ class SecurityListTestCase(TestCase): self.assertEqual(TradingControlViolation, type(exc)) exc_msg = str(ctx.exception) self.assertTrue("RestrictedListOrder" in exc_msg) - - -def add_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.txt") - with open(del_path, 'w') as f: - for sym in deletes: - f.write(sym) - f.write('\n') - add_path = os.path.join(directory, "add.txt") - with open(add_path, 'w') as f: - for sym in adds: - f.write(sym) - f.write('\n') - - -def remove_data_directory(): - directory = os.path.join( - SECURITY_LISTS_DIR, - "leveraged_etf_list/20150127/" - ) - shutil.rmtree(directory) diff --git a/zipline/finance/controls.py b/zipline/finance/controls.py index d3d77486..d38953e8 100644 --- a/zipline/finance/controls.py +++ b/zipline/finance/controls.py @@ -107,8 +107,8 @@ class RestrictedListOrder(TradingControl): def __init__(self, restricted_list): """ - restricted list can be an iterable, or - an object that implements __contains__ for dynamic + restricted list can be an iterable or a + container (implements __contains__) for dynamic restrictions. """ 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/utils/security_list.py b/zipline/utils/security_list.py index 9d242008..725d8321 100644 --- a/zipline/utils/security_list.py +++ b/zipline/utils/security_list.py @@ -7,7 +7,7 @@ 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, '..', 'security_lists') +SECURITY_LISTS_DIR = os.path.join(zipline_dir, 'resources', 'security_lists') def loopback(symbol, *args, **kwargs): @@ -25,7 +25,7 @@ class SecurityListSet(object): self._leveraged_etf = None @property - def LEVERAGED_ETF_LIST(self): + def leveraged_etf_list(self): if self._leveraged_etf is None: self._leveraged_etf = SecurityList( self.lookup_func, @@ -53,7 +53,7 @@ class SecurityList(object): self._knowledge_dates = self.make_knowledge_dates(self.data) self.current_date = current_date_func self.count = 0 - self._list = set() + self._current_set = set() def make_knowledge_dates(self, data): knowledge_dates = sorted( @@ -61,39 +61,45 @@ class SecurityList(object): return knowledge_dates def __iter__(self): - return iter(self.get_restricted_list()) + return iter(self.restricted_list) def __contains__(self, item): - rl = self.get_restricted_list() - return rl.__contains__(item) + return item in self.restricted_list - def get_restricted_list(self): + @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._list = self._cache[kd] + self._current_set = self._cache[kd] continue for effective_date, changes in self.data[kd].iteritems(): - for symbol in changes['add']: - sid = self.lookup_func( - symbol, - as_of_date=effective_date - ) - self._list.add(sid) - for symbol in changes['delete']: - sid = self.lookup_func( - symbol, - as_of_date=effective_date - ) - if sid in self._list: - self._list.remove(sid) + self.update_current( + effective_date, + changes['add'], + self._current_set.add + ) - self._cache[kd] = self._list - return self._list + 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): @@ -115,13 +121,13 @@ def load_from_directory(list_name): {add: [symbol list], 'delete': [symbol list]} """ data = {} - dir_path = SECURITY_LISTS_DIR + "/" + list_name + 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(dir_path + '/' + kd_name): + for ld_name in listdir(os.path.join(dir_path, kd_name)): ld = datetime.strptime(kd_name, DATE_FORMAT).replace( tzinfo=pytz.utc) data[kd][ld] = {} @@ -130,5 +136,6 @@ def load_from_directory(list_name): fpath = os.path.join(ld_path, fname) with open(fpath) as f: symbols = f.read().splitlines() - data[kd][ld][fname.split('.')[0]] = symbols + 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)