# # Copyright 2014 Quantopian, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. 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 from six import print_ 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 table is not 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 as err: print_(err.msg) print_("for help use --help") return 2 if __name__ == "__main__": sys.exit(main())