From eac882b773bf0bf7838debe726bf03ceb6c2ac69 Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Wed, 10 Jul 2013 14:33:14 -0400
Subject: [PATCH] ENH: Enable TALib transforms to perform on multiple stocks.

The TALib transform only supported operating on the first value
of a given batch transform panel row.

Instead of returning the one value, even if an panel with multiple
sids was provided, return a dictionary that maps stock to TALib
result.
---
 tests/test_transforms.py |  7 ++++--
 zipline/transforms/ta.py | 51 +++++++++++++++++++++++-----------------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index fb1cda7c..5aeed715 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -350,11 +350,14 @@ class TestTALIB(TestCase):
         algo = TALIBAlgorithm(talib=zipline_transforms)
         algo.run(self.source)
         # Test if computed values match those computed by pandas rolling mean.
-        talib_values = np.array(algo.talib_results[zipline_transforms[0]])
+        sid = 0
+        talib_values = np.array([x[sid] for x in
+                                 algo.talib_results[zipline_transforms[0]]])
         np.testing.assert_array_equal(talib_values,
                                       pd.rolling_mean(self.panel[0]['price'],
                                                       10).values)
-        talib_values = np.array(algo.talib_results[zipline_transforms[1]])
+        talib_values = np.array([x[sid] for x in
+                                 algo.talib_results[zipline_transforms[1]]])
         np.testing.assert_array_equal(talib_values,
                                       pd.rolling_mean(self.panel[0]['price'],
                                                       25).values)
diff --git a/zipline/transforms/ta.py b/zipline/transforms/ta.py
index a4dd2f02..82e85e72 100644
--- a/zipline/transforms/ta.py
+++ b/zipline/transforms/ta.py
@@ -30,31 +30,38 @@ def zipline_wrapper(talib_fn, key_map, data):
     else:
         req_inputs = []
 
-    # build talib_data from zipline data
-    talib_data = dict()
-    for talib_key, zipline_key in key_map.iteritems():
-        # if zipline_key is found, add it to talib_data
-        if zipline_key in data:
-            talib_data[talib_key] = data[zipline_key].values[:, 0]
-        # if zipline_key is not found and not required, add zeros
-        elif talib_key not in req_inputs:
-            talib_data[talib_key] = np.zeros(data.shape[1])
-        # if zipline key is not found and required, raise error
+    all_results = {}
+
+    for sid in data.minor_axis:
+        # build talib_data from zipline data
+        talib_data = dict()
+        for talib_key, zipline_key in key_map.iteritems():
+            # if zipline_key is found, add it to talib_data
+            if zipline_key in data:
+                talib_data[talib_key] = data[zipline_key].values[:, 0]
+            # if zipline_key is not found and not required, add zeros
+            elif talib_key not in req_inputs:
+                talib_data[talib_key] = np.zeros(data.shape[1])
+            # if zipline key is not found and required, raise error
+            else:
+                raise KeyError(
+                    'Tried to set required TA-Lib data with key '
+                    '\'{0}\' but no Zipline data is available under '
+                    'expected key \'{1}\'.'.format(
+                        talib_key, zipline_key))
+
+        # call talib
+        talib_result = talib_fn(talib_data)
+
+        # keep only the most recent result
+        if isinstance(talib_result, (list, tuple)):
+            sid_result = tuple([r[-1] for r in talib_result])
         else:
-            raise KeyError(
-                'Tried to set required TA-Lib data with key '
-                '\'{0}\' but no Zipline data is available under '
-                'expected key \'{1}\'.'.format(
-                    talib_key, zipline_key))
+            sid_result = talib_result[-1]
 
-    # call talib
-    result = talib_fn(talib_data)
+        all_results[sid] = sid_result
 
-    # keep only the most recent result
-    if isinstance(result, (list, tuple)):
-        return tuple([r[-1] for r in result])
-    else:
-        return result[-1]
+    return all_results
 
 
 def make_transform(talib_fn, name):