From 5a7fa6d8931d6f700954fe13c67bc9c435e10025 Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Mon, 2 Jul 2012 05:36:48 -0400
Subject: [PATCH] bugfixes to logging setup

---
 zipline/components/tradesimulation.py | 45 +++++++++++++++++++--------
 zipline/protocol.py                   |  4 +--
 zipline/test_algorithms.py            |  4 +--
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/zipline/components/tradesimulation.py b/zipline/components/tradesimulation.py
index 7e6f803d..94ada361 100644
--- a/zipline/components/tradesimulation.py
+++ b/zipline/components/tradesimulation.py
@@ -8,7 +8,7 @@ from zipline.core.component import Component
 from zipline.finance.trading import TransactionSimulator
 from zipline.utils.protocol_utils import  ndict
 
-from qexec.utils.log_utils import ZeroMQLogHandler, stdout_only_pipe
+from zipline.utils.log_utils import ZeroMQLogHandler, stdout_only_pipe
 
 from logbook import Logger, NestedSetup, Processor, queues
 
@@ -36,15 +36,14 @@ class TradeSimulationClient(Component):
         
         self.log_socket = log_socket
 
-        #If we have a log socket,setup context managers for capturing user prints.
+
+        #If we have a log socket,setup context manager for exporting captured
+        #print statements
         if log_socket:
-            
-            log = Logger("Print")
-            self.stdout_capture = stdout_only_pipe(log, 'user algo stdout')
-            
-            handler = queues.ZeroMQLogHandler(uri = log_socket)
-            self.zmq_out = handler.threadbound()
-        
+            self.zmq_out = ZeroMQLogHandler(uri = log_socket)
+            self.logger = Logger("Print")
+            self.stdout_capture = stdout_only_pipe #THIS IS A CLASS!
+                    
     @property
     def get_id(self):
         return str(zp.FINANCE_COMPONENT.TRADING_CLIENT)
@@ -60,9 +59,10 @@ class TradeSimulationClient(Component):
 
         # ask the algorithm to initialize, routing stdout to a zmq PUSH socket.
         if self.log_socket:
-            with self.zmq_out, self.stdout_capture:
+            with self.zmq_out.threadbound(), \
+                    self.stdout_capture(self.logger, 'Algo print capture'):
                 self.algorithm.initialize()
-        
+                
         # if we don't have a log socket, initialize anyway.
         else:
             self.algorithm.initialize()
@@ -163,7 +163,6 @@ class TradeSimulationClient(Component):
             # any fields injected here should be added to
             # LOG_EXTRA_FIELDS in zipline/protocol.py
 
-
             # try to run algo with log rerouting
             if self.log_socket:
                 def inject_event_data(record):
@@ -174,12 +173,32 @@ class TradeSimulationClient(Component):
                 log_pipeline = NestedSetup([self.zmq_out,
                                             #e.g. FileHandler(...)
                                             data_injector])
-                with log_pipeline, self.stdout_capture:
+                with log_pipeline.threadbound(), self.stdout_capture:
                     self.algorithm.handle_data(data)
             # if no log socket, just run the algo normally
             else:
                 self.algorithm.handle_data(data)
+    
+    #Testing utility for log capture.
+    def test_run_algorithm(self):
+        from zipline.utils.date_utils import epoch_now
+
+        def inject_event_data(record):
+            record.extra['algo_dt'] = epoch_now() #Mock an event.dt
+                                
+        data_injector = Processor(inject_event_data) 
+        log_pipeline = NestedSetup([self.zmq_out,
+                                    #e.g. FileHandler(...)
+                                    data_injector])
+        with log_pipeline.threadbound(), self.stdout_capture(self.logger, ''):
+            self.algorithm.handle_data('data')
+            # if no log socket, just run the algo normally
+
+
+
+        
                 
+
     def connect_order(self):
         return self.connect_push_socket(self.addresses['order_address'])
 
diff --git a/zipline/protocol.py b/zipline/protocol.py
index 89cd99aa..20d35f1c 100644
--- a/zipline/protocol.py
+++ b/zipline/protocol.py
@@ -616,7 +616,7 @@ def LOG_FRAME(payload):
        'time'      : 1199223001, #Realtime date of log creation.
        'func_name' : 'foo',
        'lineno'    : 46,
-       'message'   : 'Successfully disintegrated llama #3',
+       'msg'   : 'Successfully disintegrated llama #3',
        'level'     :  4, #Logbook enum
        'channel'   : 'MyLogger'
       }
@@ -636,7 +636,7 @@ def LOG_FRAME(payload):
         "LOG_FRAME with no channel"
     assert payload.has_key('level'),\
         "LOG_FRAME with no level"
-    assert payload.has_key('message'),\
+    assert payload.has_key('msg'),\
         "LOG_FRAME with no message"
     
     data = {}
diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py
index 059e68d5..7d33bf13 100644
--- a/zipline/test_algorithms.py
+++ b/zipline/test_algorithms.py
@@ -144,7 +144,7 @@ class TestPrintAlgorithm():
         pass
         
     def initialize(self):
-        print "Initializing"
+        print "Initializing..."
 
     def set_order(self, order_callable):
         pass
@@ -153,7 +153,7 @@ class TestPrintAlgorithm():
         pass
 
     def handle_data(self, data):
-        print "handle_data"
+        print "Handling Data..."
         pass
 
     def get_sid_filter(self):