Merge pull request #80 from quantopian/add_bars_yahoo

ENH: load_bars_from_yahoo provides OHLC. Based on Brian Cappello's code.

README.md

@@ -14,7 +14,7 @@ powering Quantopian (https://www.quantopian.com) -- a free, community-centered
 platform that allows development and real-time backtesting of trading
 algorithms in the web browser.
 
-Want to contribute? See our [open requests](https://github.com/quantopian/zipline/wiki/Contribution-Requests) 
+Want to contribute? See our [open requests](https://github.com/quantopian/zipline/wiki/Contribution-Requests)
 and our [general guidelines](https://github.com/quantopian/zipline#contributions) below.
 
 Discussion and Help
@@ -132,6 +132,7 @@ Thank you for all the help so far!
 - @snth
 - @yinhm for integrating zipline with @yinhm/datafeed
 - [Jeremiah Lowin](http://www.lowindata.com) for teaching us the nuances of Sharpe and Sortino Ratios
+- Brian Cappello
 - Quantopian Team
 
 (alert us if we've inadvertantly missed listing you here!)

tests/test_sources.py

@@ -13,11 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pandas as pd
+import pytz
 
 from unittest import TestCase
 
 import zipline.utils.factory as factory
-from zipline.sources import DataFrameSource
+from zipline.sources import DataFrameSource, DataPanelSource
 
 
 class TestDataFrameSource(TestCase):
@@ -52,3 +53,22 @@ class TestDataFrameSource(TestCase):
             self.assertEquals(event['sid'], 0)
             self.assertTrue(isinstance(event['volume'], int))
             self.assertTrue(isinstance(event['arbitrary'], float))
+
+    def test_yahoo_bars_to_panel_source(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 = factory.load_bars_from_yahoo(stocks=stocks,
+                                            start=start,
+                                            end=end)
+
+        source = DataPanelSource(data)
+        for event in source:
+            self.assertTrue('sid' in event)
+            self.assertTrue('open' in event)
+            self.assertTrue('high' in event)
+            self.assertTrue('low' in event)
+            self.assertTrue('close' in event)
+            self.assertTrue('volume' in event)
+            self.assertTrue('price' in event)
+            self.assertTrue(isinstance(event['volume'], (int, long)))

tests/test_utils.py

@@ -14,7 +14,8 @@
 # limitations under the License.
 
 from unittest import TestCase
-from zipline.utils.factory import load_from_yahoo
+from zipline.utils.factory import (load_from_yahoo,
+                                   load_bars_from_yahoo)
 import pandas as pd
 import pytz
 import numpy as np
@@ -36,3 +37,22 @@ class TestFactory(TestCase):
             AssertionError, load_from_yahoo, stocks=stocks,
             start=end, end=start
         )
+
+    def test_load_bars_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_bars_from_yahoo(stocks=stocks, start=start, end=end)
+
+        assert data.major_axis[0] == pd.Timestamp('1993-01-04 00:00:00+0000')
+        assert data.major_axis[-1] == pd.Timestamp('2001-12-31 00:00:00+0000')
+        for stock in stocks:
+            assert stock in data.items
+
+        for ohlc in ['open', 'high', 'low', 'close', 'volume', 'price']:
+            assert ohlc in data.minor_axis
+
+        np.testing.assert_raises(
+            AssertionError, load_bars_from_yahoo, stocks=stocks,
+            start=end, end=start
+        )

zipline/sources/data_frame_source.py

@@ -1,3 +1,4 @@
+
 #
 # Copyright 2012 Quantopian, Inc.
 #

zipline/utils/factory.py

@@ -357,7 +357,7 @@ def create_test_panel_source(trading_calendar=None):
     return DataPanelSource(panel), panel
 
 
-def load_from_yahoo(indexes=None, stocks=None, start=None, end=None):
+def _load_raw_yahoo_data(indexes=None, stocks=None, start=None, end=None):
     """Load closing prices from yahoo finance.
 
     :Optional:
@@ -398,8 +398,21 @@ def load_from_yahoo(indexes=None, stocks=None, start=None, end=None):
         print name
         stkd = DataReader(ticker, 'yahoo', start, end).sort_index()
         data[name] = stkd
+    return data
 
-    df = pd.DataFrame({key: d['Close'] for key, d in data.iteritems()})
+
+def load_from_yahoo(indexes=None, stocks=None, start=None, end=None):
+    data = _load_raw_yahoo_data(indexes, stocks, start, end)
+    df = pd.DataFrame({key: d['Adj Close'] for key, d in data.iteritems()})
     df.index = df.index.tz_localize(pytz.utc)
-
     return df
+
+
+def load_bars_from_yahoo(indexes=None, stocks=None, start=None, end=None):
+    data = _load_raw_yahoo_data(indexes, stocks, start, end)
+    panel = pd.Panel(data)
+    # Rename columns
+    panel.minor_axis = ['open', 'high', 'low', 'close', 'volume', 'price']
+    panel.major_axis = panel.major_axis.tz_localize(pytz.utc)
+
+    return panel