From 7d5194ec2c9dfecf00148364a25bdcff956e6824 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Mon, 19 Aug 2013 16:46:11 -0400 Subject: [PATCH] ENH: Include TALib output names in ta transform results. For TALib functions like MACD that have output names, return a DataFrame that for which the columns are the output names of the function. So that when using a TALib function, the algorithm doesn't need to know the index position of the desired result, in favor of using the name of the result. e.g. ``` macd_result['AAPL'][0] ``` becomes, ``` macd_result['AAPL']['macd'] ``` and ``` macd_result['AAPL'][1] ``` becomes, ``` macd_result['AAPL']['macdsignal'] ``` Also, change return type of functions that return floats from a dictionary to a Series, so that the function is always returning a pandas type. --- zipline/transforms/ta.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/zipline/transforms/ta.py b/zipline/transforms/ta.py index 0c83df94..54752fb8 100644 --- a/zipline/transforms/ta.py +++ b/zipline/transforms/ta.py @@ -31,7 +31,16 @@ def zipline_wrapper(talib_fn, key_map, data): else: req_inputs = [] - all_results = {} + # If there are multiple output names then the results are named, + # if there is only one output name, it usually 'real' is best represented + # by a float. + # Use a DataFrame to map sid to named values, and a Series map sid + # to floats. + if len(talib_fn.output_names) > 1: + all_results = pd.DataFrame(index=talib_fn.output_names, + columns=data.minor_axis) + else: + all_results = pd.Series(index=data.minor_axis) for sid in data.minor_axis: # build talib_data from zipline data