From f81addb7df9979430338de9edfca88f9046190f3 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Mon, 3 Dec 2012 11:48:45 -0500
Subject: [PATCH] ENH: batch_transform now adds arbitrary keys to datapanel.

---
 tests/test_transforms.py    | 12 ++++++++++++
 zipline/test_algorithms.py  | 21 +++++++++++++++++++++
 zipline/transforms/utils.py | 22 +++++++++++++++-------
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index c96d5e31..ae340ca8 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -336,6 +336,18 @@ class TestBatchTransform(TestCase):
             "Sixth iteration should not be None"
         )
 
+        # Test whether arbitrary fields can be added to datapanel
+        field = algo.history_return_arbitrary_fields[-1]
+        self.assertTrue(
+            'arbitrary' in field.items,
+            'datapanel should contain column arbitrary'
+        )
+
+        self.assertTrue(all(
+            field['arbitrary'].values.flatten() == ['test'] * 8),
+            'arbitrary dataframe should contain only "test"'
+        )
+
         # test overloaded class
         for test_history in [algo.history_return_price_class,
                              algo.history_return_price_decorator]:
diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py
index 93fd09aa..8fa80b22 100644
--- a/zipline/test_algorithms.py
+++ b/zipline/test_algorithms.py
@@ -71,6 +71,8 @@ The algorithm must expose methods:
     and trade events.
 
 """
+from copy import deepcopy
+
 from zipline.algorithm import TradingAlgorithm
 from zipline.finance.slippage import FixedSlippage
 
@@ -248,6 +250,11 @@ def return_args_batch_decorator(data, *args, **kwargs):
     return args, kwargs
 
 
+@batch_transform
+def return_data(data, *args, **kwargs):
+    return data
+
+
 class BatchTransformAlgorithm(TradingAlgorithm):
     def initialize(self, *args, **kwargs):
         self.refresh_period = kwargs.pop('refresh_period', 2)
@@ -261,6 +268,7 @@ class BatchTransformAlgorithm(TradingAlgorithm):
         self.history_return_args = []
         self.history_return_price_market_aware = []
         self.history_return_more_days_than_refresh = []
+        self.history_return_arbitrary_fields = []
 
         self.return_price_class = ReturnPriceBatchTransform(
             market_aware=False,
@@ -292,6 +300,12 @@ class BatchTransformAlgorithm(TradingAlgorithm):
             window_length=3
         )
 
+        self.return_arbitrary_fields = return_data(
+            market_aware=True,
+            refresh_period=1,
+            window_length=3
+        )
+
         self.set_slippage(FixedSlippage())
 
     def handle_data(self, data):
@@ -307,6 +321,13 @@ class BatchTransformAlgorithm(TradingAlgorithm):
         self.history_return_more_days_than_refresh.append(
             self.return_price_more_days_than_refresh.handle_data(data))
 
+        new_data = deepcopy(data)
+        for sid in new_data:
+            new_data[sid]['arbitrary'] = 'test'
+
+        self.history_return_arbitrary_fields.append(
+            self.return_arbitrary_fields.handle_data(new_data))
+
 
 class SetPortfolioAlgorithm(TradingAlgorithm):
     """
diff --git a/zipline/transforms/utils.py b/zipline/transforms/utils.py
index 04157158..1fd965c5 100644
--- a/zipline/transforms/utils.py
+++ b/zipline/transforms/utils.py
@@ -118,7 +118,6 @@ class StatefulTransform(object):
         # messages should only manipulate copies.
         log.info('Running StatefulTransform [%s]' % self.get_hash())
         for message in stream_in:
-
             # allow upstream generators to yield None to avoid
             # blocking.
             if message is None:
@@ -233,7 +232,6 @@ class EventWindow(object):
 
     def update(self, event):
         self.assert_well_formed(event)
-
         # Add new event and increment totals.
         self.ticks.append(deepcopy(event))
 
@@ -370,6 +368,8 @@ class BatchTransform(EventWindow):
         self.updated = False
         self.data = None
 
+        self.field_names = None
+
     def handle_data(self, data, *args, **kwargs):
         """
         New method to handle a data frame as sent to the algorithm's
@@ -399,6 +399,16 @@ class BatchTransform(EventWindow):
             self.last_dt = event.dt
             return
 
+        # extract field names from sids (price, volume etc), make sure
+        # every sid has the same fields.
+        sid_keys = [set(sid.keys()) for sid in event.data.itervalues()]
+        assert sid_keys[0] == set.union(*sid_keys),\
+            "Each sid must have the same keys."
+        if self.field_names is None:
+            unwanted_fields = set(['portfolio', 'sid', 'dt', 'type',
+                                   'datetime'])
+            self.field_names = sid_keys[0] - unwanted_fields
+
         # update trading day counters
         if self.last_dt.day != event.dt.day:
             self.last_dt = event.dt
@@ -436,11 +446,9 @@ class BatchTransform(EventWindow):
         # self.ticks contains ndicts with data, dt keys.
         # event parameter is an ndict with data, dt keys.
         fields = {}
-        for field_name in ['price', 'volume']:
+
+        for field_name in self.field_names:
             sids = self.ticks[0].data.keys()
-            # Skip non-existant fields
-            if field_name not in self.ticks[0].data[sids[0]]:
-                continue
 
             values_per_sid = {}
 
@@ -452,7 +460,7 @@ class BatchTransform(EventWindow):
 
                 # concatenate different sids into one df
                 df = pd.DataFrame.from_dict(values_per_sid)
-                 # Fills in gaps of missing data during transform of multiple
+                # Fills in gaps of missing data during transform of multiple
                 # stocks.
                 # e.g. we may be missing minute data because of illiquidity
                 # of one stock