From 956107a84661f26c30c65ee3a4f051941cd4507e Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Mon, 30 Sep 2013 11:25:59 -0400 Subject: [PATCH] MAINT: Use pandas instead of msgpack for benchmarks and treasuries. Instead of writing our own serialization using msgpack, leverage the csv serialization provided by pandas. Also, lessens the need for msgpack and functions in date_utils. --- zipline/data/loader.py | 121 +++++++++++++++++++---------------------- 1 file changed, 55 insertions(+), 66 deletions(-) diff --git a/zipline/data/loader.py b/zipline/data/loader.py index ad0db0e4..741ea3f8 100644 --- a/zipline/data/loader.py +++ b/zipline/data/loader.py @@ -16,20 +16,19 @@ import os from os.path import expanduser -import msgpack from collections import OrderedDict from datetime import timedelta import logbook +import pandas as pd + from . treasuries import get_treasury_data from . import benchmarks from . benchmarks import get_benchmark_returns from zipline.protocol import DailyReturn -from zipline.utils.date_utils import tuple_to_date from zipline.utils.tradingcalendar import trading_days -from operator import attrgetter logger = logbook.Logger('Loader') @@ -60,18 +59,17 @@ def dump_treasury_curves(): Puts source treasury and data into zipline. """ - tr_data = [] + tr_data = {} for curve in get_treasury_data(): - date_as_tuple = curve['date'].timetuple()[0:6] + \ - (curve['date'].microsecond,) # Not ideal but massaging data into expected format - del curve['date'] - tr = (date_as_tuple, curve) - tr_data.append(tr) + tr_data[curve['date']] = curve - with get_datafile('treasury_curves.msgpack', mode='wb') as tr_fp: - tr_fp.write(msgpack.dumps(tr_data)) + curves = pd.DataFrame(tr_data).T + + datafile = get_datafile('treasury_curves.csv', mode='wb') + curves.to_csv(datafile) + datafile.close() def dump_benchmarks(symbol): @@ -82,14 +80,14 @@ def dump_benchmarks(symbol): """ benchmark_data = [] for daily_return in get_benchmark_returns(symbol): - date_as_tuple = daily_return.date.timetuple()[0:6] + \ - (daily_return.date.microsecond,) # Not ideal but massaging data into expected format - benchmark = (date_as_tuple, daily_return.returns) + benchmark = (daily_return.date, daily_return.returns) benchmark_data.append(benchmark) - with get_datafile(get_benchmark_filename(symbol), mode='wb') as bmark_fp: - bmark_fp.write(msgpack.dumps(benchmark_data)) + datafile = get_datafile(get_benchmark_filename(symbol), mode='wb') + benchmark_returns = pd.Series(dict(benchmark_data)) + benchmark_returns.to_csv(datafile) + datafile.close() def update_treasury_curves(last_date): @@ -100,22 +98,18 @@ def update_treasury_curves(last_date): Puts source treasury and data into zipline. """ - tr_data = [] - with get_datafile('treasury_curves.msgpack', mode='rb') as tr_fp: - tr_list = msgpack.loads(tr_fp.read()) - for packed_date, curve in tr_list: - tr_data.append((packed_date, curve)) + datafile = get_datafile('treasury_curves.csv', mode='rb') + curves = pd.DataFrame.from_csv(datafile).T + datafile.close() for curve in get_treasury_data(): - date_as_tuple = curve['date'].timetuple()[0:6] + \ - (curve['date'].microsecond,) - # Not ideal but massaging data into expected format - del curve['date'] - tr = (date_as_tuple, curve) - tr_data.append(tr) + curves[curve['date']] = curve - with get_datafile('treasury_curves.msgpack', mode='wb') as tr_fp: - tr_fp.write(msgpack.dumps(tr_data)) + updated_curves = curves.T + + datafile = get_datafile('treasury_curves.csv', mode='wb') + updated_curves.to_csv(datafile) + datafile.close() def update_benchmarks(symbol, last_date): @@ -126,30 +120,27 @@ def update_benchmarks(symbol, last_date): Puts source benchmark into zipline. """ - benchmark_data = [] - with get_datafile(get_benchmark_filename(symbol), mode='rb') as bmark_fp: - bm_list = msgpack.loads(bmark_fp.read()) - for packed_date, returns in bm_list: - benchmark_data.append((packed_date, returns)) + datafile = get_datafile(get_benchmark_filename(symbol), mode='rb') + saved_benchmarks = pd.Series.from_csv(datafile) + datafile.close() try: start = last_date + timedelta(days=1) for daily_return in get_benchmark_returns(symbol, start_date=start): - date_as_tuple = daily_return.date.timetuple()[0:6] + \ - (daily_return.date.microsecond,) # Not ideal but massaging data into expected format - benchmark = (date_as_tuple, daily_return.returns) - benchmark_data.append(benchmark) + benchmark = pd.Series({daily_return.date: daily_return.returns}) + saved_benchmarks.append(benchmark) - with get_datafile( - get_benchmark_filename(symbol), mode='wb') as bmark_fp: - bmark_fp.write(msgpack.dumps(benchmark_data)) + datafile = get_datafile(get_benchmark_filename(symbol), mode='wb') + saved_benchmarks.to_csv(datafile) + datafile.close() except benchmarks.BenchmarkDataNotFoundError as exc: logger.warn(exc) + return saved_benchmarks def get_benchmark_filename(symbol): - return "%s_benchmark.msgpack" % symbol + return "%s_benchmark.csv" % symbol def load_market_data(bm_symbol='^GSPC'): @@ -157,69 +148,67 @@ def load_market_data(bm_symbol='^GSPC'): fp_bm = get_datafile(get_benchmark_filename(bm_symbol), "rb") except IOError: print(""" -data msgpacks aren't distributed with source. +data files aren't distributed with source. Fetching data from Yahoo Finance. """).strip() dump_benchmarks(bm_symbol) fp_bm = get_datafile(get_benchmark_filename(bm_symbol), "rb") - bm_list = msgpack.loads(fp_bm.read()) + saved_benchmarks = pd.Series.from_csv(fp_bm) + fp_bm.close() # Find the offset of the last date for which we have trading data in our # list of valid trading days - last_bm_date = tuple_to_date(bm_list[-1][0]) + last_bm_date = saved_benchmarks.index[-1] last_bm_date_offset = trading_days.searchsorted( last_bm_date.strftime('%Y/%m/%d')) # If more than 1 trading days has elapsed since the last day where # we have data,then we need to update if len(trading_days) - last_bm_date_offset > 1: - update_benchmarks(bm_symbol, last_bm_date) - fp_bm = get_datafile(get_benchmark_filename(bm_symbol), "rb") - bm_list = msgpack.loads(fp_bm.read()) + benchmark_returns = update_benchmarks(bm_symbol, last_bm_date) + else: + benchmark_returns = saved_benchmarks + + benchmark_returns = benchmark_returns.tz_localize('UTC') bm_returns = [] - for packed_date, returns in bm_list: - event_dt = tuple_to_date(packed_date) - - daily_return = DailyReturn(date=event_dt, returns=returns) + for dt, returns in benchmark_returns.iterkv(): + daily_return = DailyReturn(date=dt, returns=returns) bm_returns.append(daily_return) - fp_bm.close() - - bm_returns = sorted(bm_returns, key=attrgetter('date')) - try: - fp_tr = get_datafile('treasury_curves.msgpack', "rb") + fp_tr = get_datafile('treasury_curves.csv', "rb") except IOError: print(""" data msgpacks aren't distributed with source. Fetching data from data.treasury.gov """).strip() dump_treasury_curves() - fp_tr = get_datafile('treasury_curves.msgpack', "rb") + fp_tr = get_datafile('treasury_curves.csv', "rb") - tr_list = msgpack.loads(fp_tr.read()) + saved_curves = pd.DataFrame.from_csv(fp_tr) # Find the offset of the last date for which we have trading data in our # list of valid trading days - last_tr_date = tuple_to_date(tr_list[-1][0]) + last_tr_date = saved_curves.index[-1] last_tr_date_offset = trading_days.searchsorted( last_tr_date.strftime('%Y/%m/%d')) # If more than 1 trading days has elapsed since the last day where # we have data,then we need to update if len(trading_days) - last_tr_date_offset > 1: - update_treasury_curves(last_tr_date) - fp_tr = get_datafile('treasury_curves.msgpack', "rb") - tr_list = msgpack.loads(fp_tr.read()) + treasury_curves = update_treasury_curves(last_tr_date) + else: + treasury_curves = saved_curves + + treasury_curves = treasury_curves.tz_localize('UTC') tr_curves = {} - for packed_date, curve in tr_list: - tr_dt = tuple_to_date(packed_date) + for tr_dt, curve in treasury_curves.T.iterkv(): # tr_dt = tr_dt.replace(hour=0, minute=0, second=0, microsecond=0, # tzinfo=pytz.utc) - tr_curves[tr_dt] = curve + tr_curves[tr_dt] = curve.to_dict() fp_tr.close()