diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..b5fafa7d --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,38 @@ +# +# Copyright 2012 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. + +from unittest2 import TestCase +from zipline.utils.factory import load_from_yahoo +import pandas as pd +import pytz +import numpy as np + + +class TestFactory(TestCase): + def test_load_from_yahoo(self): + stocks = ['AAPL', 'GE'] + start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc) + end = pd.datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc) + data = load_from_yahoo(stocks=stocks, start=start, end=end) + + assert data.index[0] == pd.Timestamp('1993-01-04 00:00:00+0000') + assert data.index[-1] == pd.Timestamp('2001-12-31 00:00:00+0000') + for stock in stocks: + assert stock in data.columns + + np.testing.assert_raises( + AssertionError, load_from_yahoo, stocks=stocks, + start=end, end=start + ) diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index f2f4ed67..86a19183 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -22,8 +22,10 @@ import msgpack import random from os.path import join, abspath, dirname from operator import attrgetter +from collections import OrderedDict import pandas as pd +from pandas.io.data import DataReader import numpy as np from datetime import datetime, timedelta @@ -276,3 +278,51 @@ def create_test_df_source(): df = pd.DataFrame(x, index=index, columns=[0, 1]) return DataFrameSource(df), df + + +def load_from_yahoo(indexes=None, stocks=None, start=None, end=None): + """Load closing prices from yahoo finance. + + :Optional: + indexes : dict (Default: {'SPX': '^GSPC'}) + Financial indexes to load. + stocks : list (Default: ['AAPL', 'GE', 'IBM', 'MSFT', + 'XOM', 'AA', 'JNJ', 'PEP', 'KO']) + Stock closing prices to load. + start : datetime (Default: datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)) + Retrieve prices from start date on. + end : datetime (Default: datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc)) + Retrieve prices until end date. + + :Note: + This is based on code presented in a talk by Wes McKinney: + http://wesmckinney.com/files/20111017/notebook_output.pdf + """ + + if indexes is None: + indexes = {'SPX': '^GSPC'} + if stocks is None: + stocks = ['AAPL', 'GE', 'IBM', 'MSFT', 'XOM', 'AA', 'JNJ', 'PEP', 'KO'] + if start is None: + start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc) + if end is None: + end = pd.datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc) + + assert start < end, "start date is later than end date." + + data = OrderedDict() + + for stock in stocks: + print stock + stkd = DataReader(stock, 'yahoo', start, end).sort_index() + data[stock] = stkd + + for name, ticker in indexes.iteritems(): + print name + stkd = DataReader(ticker, 'yahoo', start, end).sort_index() + data[name] = stkd + + df = pd.DataFrame({key: d['Close'] for key, d in data.iteritems()}) + df.index = df.index.tz_localize(pytz.utc) + + return df