Merge pull request #478 from quantopian/dnt_list

ENH: no order guard, security lists

MANIFEST.in

@@ -1,2 +1,4 @@
 include README.md
 include LICENSE
+
+recursive-include zipline/resources *.*

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=[

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.

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)

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

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

zipline/resources/security_lists/leveraged_etf_list/20150126/20120913/add

@@ -0,0 +1,2 @@
+ROSA
+SFSA

zipline/resources/security_lists/leveraged_etf_list/20150126/20120919/add

@@ -0,0 +1 @@
+RTSA

zipline/resources/security_lists/leveraged_etf_list/20150126/20121012/add

@@ -0,0 +1 @@
+VZZB

zipline/resources/security_lists/leveraged_etf_list/20150126/20130605/add

@@ -0,0 +1,2 @@
+LSKY
+SSDL

zipline/resources/security_lists/leveraged_etf_list/20150126/20130916/add

@@ -0,0 +1 @@
+MFSA

zipline/resources/security_lists/leveraged_etf_list/20150126/20131002/add

@@ -0,0 +1,2 @@
+DSTJ
+DSXJ

zipline/resources/security_lists/leveraged_etf_list/20150126/20131009/add

@@ -0,0 +1 @@
+JFT

zipline/resources/security_lists/leveraged_etf_list/20150126/20131121/add

@@ -0,0 +1,4 @@
+FSA
+FSE
+FSG
+FSU

zipline/resources/security_lists/leveraged_etf_list/20150126/20131227/add

@@ -0,0 +1,2 @@
+IPLT
+LPLT

zipline/resources/security_lists/leveraged_etf_list/20150126/20140410/add

@@ -0,0 +1 @@
+BXDB

zipline/resources/security_lists/leveraged_etf_list/20150126/20140923/add

@@ -0,0 +1 @@
+BRZS

zipline/resources/security_lists/leveraged_etf_list/20150126/20141119/add

@@ -0,0 +1,2 @@
+BXUB
+BXUC

zipline/resources/security_lists/leveraged_etf_list/20150126/20141226/add

@@ -0,0 +1 @@
+BARS

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

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

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

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)