From 9e1a5c11cbfcc5011ccd9065926a86bb387e1f30 Mon Sep 17 00:00:00 2001 From: scottsanderson Date: Wed, 1 Aug 2012 18:56:29 -0400 Subject: [PATCH] fix syntax in transform --- zipline/gens/transform.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 2adc44e8..9d90aeb8 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -79,11 +79,21 @@ def stateful_transform(stream_in, tnfm_class, *args, **kwargs): out_message.tnfm_value = tnfm_value yield out_message - # Special logic for TransactionSimulator. This is ugly but I - # want to get to testing faster. Should be refactored later - # to something that doesn't make Scott cry. - elif tnfm_class.__name__ == 'TransactionSimulator' + # Special logic for TransactionSimulator and + # PerformanceTracker. This is ugly but I want to get to + # testing faster. Should be refactored later to something + # that doesn't make Scott cry. + elif tnfm_class.__name__ == 'TransactionSimulator': + out_message = message_copy + out_message.txn = tnfm_value + yield out_message + elif tnfm_class.__name__ == 'PerformanceTracker': + out_message = message_copy + del out_message['txn'] + out_message.portfolio = tnfm_value + yield out_message + # Otherwise send tnfm_id, tnfm_value, and the message date. else: out_message = ndict()