ENH: make BcolzMinuteBarWriter.write take iterable

Updates the BcolzMinuteBarWriter.write api to allow users to pass their
data as a stream instead of requiring that they loop over their data
externally. This matches the API presented by BcolzDailyBarWriter.

zipline/data/minute_bars.py

@@ -29,6 +29,7 @@ from zipline.data._minute_bar_internal import (
 )
 
 from zipline.gens.sim_engine import NANOS_IN_MINUTE
+from zipline.utils.cli import maybe_show_progress
 from zipline.utils.memoize import lazyval
 
 US_EQUITIES_MINUTES_PER_DAY = 390
@@ -441,7 +442,36 @@ class BcolzMinuteBarWriter(object):
         assert new_last_date == date, "new_last_date={0} != date={1}".format(
             new_last_date, date)
 
-    def write(self, sid, df):
+    def write(self, data, show_progress=False):
+        """Write a stream of minute data.
+
+        Parameters
+        ----------
+        data : iterable[(int, pd.DataFrame)]
+            The data to write. Each element should be a tuple of sid, data
+            where data has the following format:
+            columns : ('open', 'high', 'low', 'close', 'volume')
+                open : float64
+                high : float64
+                low  : float64
+                close : float64
+                volume : float64|int64
+            index : DatetimeIndex of market minutes.
+        show_progress : bool, optional
+            Whether or not to show a progress bar while writing.
+        """
+        ctx = maybe_show_progress(
+            data,
+            show_progress=show_progress,
+            item_show_func=lambda e: e if e is None else str(e[0]),
+            label="Merging minute equity files:",
+        )
+        write_sid = self.write_sid
+        with ctx as it:
+            for e in it:
+                write_sid(*e)
+
+    def write_sid(self, sid, df):
         """
         Write the OHLCV data for the given sid.
         If there is no bcolz ctable yet created for the sid, create it.

zipline/data/us_equity_pricing.py

@@ -209,7 +209,7 @@ class BcolzDailyBarWriter(object):
 
     @property
     def progress_bar_message(self):
-        return "Merging asset files:"
+        return "Merging daily equity files:"
 
     def progress_bar_item_show_func(self, value):
         return value if value is None else str(value[0])
@@ -229,9 +229,9 @@ class BcolzDailyBarWriter(object):
             The assets that should be in ``data``. If this is provided
             we will check ``data`` against the assets and provide better
             progress information.
-        show_progress : bool
+        show_progress : bool, optional
             Whether or not to show a progress bar while writing.
-        invalid_data_behavior : {'warn', 'raise', 'ignore'}
+        invalid_data_behavior : {'warn', 'raise', 'ignore'}, optional
             What to do when data is encountered that is outside the range of
             a uint32.
 

zipline/testing/core.py

@@ -349,7 +349,7 @@ def make_trade_data_for_asset_info(dates,
         )
 
         if writer:
-            writer.write(sid, df)
+            writer.write_sid(sid, df)
 
         trade_data[sid] = df
 
@@ -424,8 +424,8 @@ def write_minute_data(env, tempdir, minutes, sids):
 
 def create_minute_bar_data(minutes, sids):
     length = len(minutes)
-    return {
-        sid: pd.DataFrame(
+    for sid_idx, sid in enumerate(sids):
+        yield sid, pd.DataFrame(
             {
                 'open': np.arange(length) + 10 + sid_idx,
                 'high': np.arange(length) + 15 + sid_idx,
@@ -435,8 +435,6 @@ def create_minute_bar_data(minutes, sids):
             },
             index=minutes,
         )
-        for sid_idx, sid in enumerate(sids)
-    }
 
 
 def create_daily_bar_data(trading_days, sids):
@@ -492,20 +490,17 @@ def create_data_portal(env, tempdir, sim_params, sids, adjustment_reader=None):
         )
 
 
-def write_bcolz_minute_data(env, days, path, df_dict):
+def write_bcolz_minute_data(env, days, path, data):
     market_opens = env.open_and_closes.market_open.loc[days]
     market_closes = env.open_and_closes.market_close.loc[days]
 
-    writer = BcolzMinuteBarWriter(
+    BcolzMinuteBarWriter(
         days[0],
         path,
         market_opens,
         market_closes,
         US_EQUITIES_MINUTES_PER_DAY
-    )
-
-    for sid, df in iteritems(df_dict):
-        writer.write(sid, df)
+    ).write(data)
 
 
 def create_minute_df_for_asset(env,

zipline/testing/fixtures.py

@@ -6,7 +6,7 @@ from contextlib2 import ExitStack
 from logbook import NullHandler, Logger
 from nose_parameterized import parameterized
 from pandas.util.testing import assert_series_equal
-from six import with_metaclass, iteritems
+from six import with_metaclass
 from toolz import flip
 import numpy as np
 import pandas as pd
@@ -681,7 +681,7 @@ class WithBcolzMinuteBarReader(WithTradingEnvironment, WithTmpDir):
 
     Methods
     -------
-    make_minute_bar_data() -> dict[int -> pd.DataFrame]
+    make_minute_bar_data() -> iterable[(int, pd.DataFrame)]
         A class method that returns a dict mapping sid to dataframe
         which will be written to the bcolz files that the class's
         ``BcolzMinuteBarReader`` will read from. By default this creates
@@ -734,9 +734,7 @@ class WithBcolzMinuteBarReader(WithTradingEnvironment, WithTmpDir):
             cls.env.open_and_closes.market_close.loc[days],
             US_EQUITIES_MINUTES_PER_DAY
         )
-        cls.bcolz_minute_bar_data = cls.make_minute_bar_data()
-        for sid, df in iteritems(cls.bcolz_minute_bar_data):
-            writer.write(sid, df)
+        writer.write(cls.make_minute_bar_data())
 
         cls.bcolz_minute_bar_reader = BcolzMinuteBarReader(p)