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

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