Merge pull request #2 from quantopian/yahoo_finance_loader

Added yahoo finance loader.

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
+        )

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