catalyst/zipline/utils/data_source_tables_gen.py

#!/usr/bin/env python
import sys
import getopt
import traceback
import numpy as np
import pandas as pd
import datetime
import logging
import tables
import gzip
import glob
import os
import random
import csv
import time
FORMAT = "%(asctime)-15s -8s %(message)s"

logging.basicConfig(format=FORMAT, level=logging.INFO)


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


OHLCTableDescription = {'sid': tables.StringCol(14, pos=2),
                        'dt': tables.Int64Col(pos=1),
                        'open': tables.Float64Col(dflt=np.NaN, pos=3),
                        'high': tables.Float64Col(dflt=np.NaN, pos=4),
                        'low': tables.Float64Col(dflt=np.NaN, pos=5),
                        'close': tables.Float64Col(dflt=np.NaN, pos=6),
                        "volume": tables.Int64Col(dflt=0, pos=7)}


def process_line(line):
    dt = np.datetime64(line["dt"]).astype(np.int64)
    sid = line["sid"]
    open_p = float(line["open"])
    high_p = float(line["high"])
    low_p = float(line["low"])
    close_p = float(line["close"])
    volume = int(line["volume"])
    return (dt, sid, open_p, high_p, low_p, close_p, volume)


def parse_csv(csv_reader):
    previous_date = None
    data = []
    dtype = [('dt', 'int64'), ('sid', '|S14'), ('open', float),
             ('high', float), ('low', float), ('close', float),
             ('volume', int)]
    for line in csv_reader:
        row = process_line(line)
        current_date = line["dt"][:10].replace("-", "")
        if previous_date and previous_date != current_date:
            rows = np.array(data, dtype=dtype).view(np.recarray)
            yield current_date, rows
            data = []
        data.append(row)
        previous_date = current_date


def merge_all_files_into_pytables(file_dir, file_out):
    """
    process each file into pytables
    """
    start = None
    start = datetime.datetime.now()
    out_h5 = tables.openFile(file_out,
                             mode="w",
                             title="bars",
                             filters=tables.Filters(complevel=9,
                                                    complib='zlib'))
    table = None
    for file_in in glob.glob(file_dir + "/*.gz"):
        gzip_file = gzip.open(file_in)
        expected_header = ["dt", "sid", "open", "high", "low", "close",
                           "volume"]
        csv_reader = csv.DictReader(gzip_file)
        header = csv_reader.fieldnames
        if header != expected_header:
            logging.warn("expected header %s\n" % (expected_header))
            logging.warn("header_found %s" % (header))
            return

        for current_date, rows in parse_csv(csv_reader):
            table = out_h5.createTable("/TD", "date_" + current_date,
                                       OHLCTableDescription,
                                       expectedrows=len(rows),
                                       createparents=True)
            table.append(rows)
            table.flush()
        if not table is None:
            table.flush()
    end = datetime.datetime.now()
    diff = (end - start).seconds
    logging.debug("finished  it took %d." % (diff))


def create_fake_csv(file_in):
    fields = ["dt", "sid", "open", "high", "low", "close", "volume"]
    gzip_file = gzip.open(file_in, "w")
    dict_writer = csv.DictWriter(gzip_file, fieldnames=fields)
    current_dt = datetime.date.today() - datetime.timedelta(days=2)
    current_dt = pd.Timestamp(current_dt).replace(hour=9)
    current_dt = current_dt.replace(minute=30)
    end_time = pd.Timestamp(datetime.date.today())
    end_time = end_time.replace(hour=16)
    last_price = 10.0
    while current_dt < end_time:
        row = {}
        row["dt"] = current_dt
        row["sid"] = "test"
        last_price += random.randint(-20, 100) / 10000.0
        row["close"] = last_price
        row["open"] = last_price - 0.01
        row["low"] = last_price - 0.02
        row["high"] = last_price + 0.02
        row["volume"] = random.randint(10, 1000) * 10
        dict_writer.writerow(row)
        current_dt += datetime.timedelta(minutes=1)
        if current_dt.hour > 16:
            current_dt += datetime.timedelta(days=1)
            current_dt = current_dt.replace(hour=9)
            current_dt = current_dt.replace(minute=30)
    gzip_file.close()


def main(argv=None):
    """
    This script cleans minute bars into pytables file
    data_source_tables_gen.py
    [--tz_in] sets time zone of data only reasonably fast way to use
    time.tzset()
    [--dir_in] iterates through directory provided of csv files in gzip form
    in form:
    dt, sid, open, high, low, close, volume
    2012-01-01T12:30:30,1234HT,1, 2,3,4.0
    [--fake_csv] creates a fake sample csv to iterate through
    [--file_out] determines output file
    """
    if argv is None:
        argv = sys.argv
    try:
        dir_in = None
        file_out = "./all.h5"
        fake_csv = None
        try:
            opts, args = getopt.getopt(argv[1:], "hdft",
                                       ["help",
                                        "dir_in=",
                                        "debug",
                                        "tz_in=",
                                        "fake_csv=",
                                        "file_out="])
        except getopt.error as msg:
            raise Usage(msg)
        for opt, value in opts:
            if opt in ("--help", "-h"):
                print main.__doc__
            if opt in ("-d", "--debug"):
                logging.basicConfig(format=FORMAT,
                                    level=logging.DEBUG)
            if opt in ("-d", "--dir_in"):
                dir_in = value
            if opt in ("-o", "--file_out"):
                file_out = value
            if opt in ("--fake_csv"):
                fake_csv = value
            if opt in ("--tz_in"):
                os.environ['TZ'] = value
                time.tzset()
        try:
            if dir_in:
                merge_all_files_into_pytables(dir_in, file_out)
            if fake_csv:
                create_fake_csv(fake_csv)
        except Exception:
            error = "An unhandled error occured in the"
            error += "data_source_tables_gen.py script."
            error += "\n\nTraceback:\n"
            error += '-' * 70 + "\n"
            error += "".join(traceback.format_tb(sys.exc_info()[2]))
            error += repr(sys.exc_info()[1]) + "\n"
            error += str(sys.exc_info()[1]) + "\n"
            error += '-' * 70 + "\n"
            print error
    except Usage, err:
        print >>sys.stderr, err.msg
        print >>sys.stderr, "for help use --help"
        return 2

if __name__ == "__main__":
    sys.exit(main())