catalyst/zipline/utils/security_list.py

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, '..', '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._list = 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.get_restricted_list())

    def __contains__(self, item):
        rl = self.get_restricted_list()
        return rl.__contains__(item)

    def get_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]
                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._cache[kd] = self._list
        return self._list


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 = 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):
            ld = datetime.strptime(kd_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.split('.')[0]] = symbols
    return data