capturing stdout from user's algo and piping it to a zmq socket

zipline/components/tradesimulation.py

@@ -1,4 +1,5 @@
 import logging
+import logbook
 import datetime
 
 import zipline.protocol as zp
@@ -7,12 +8,16 @@ import zipline.finance.performance as perf
 from zipline.core.component import Component
 from zipline.finance.trading import TransactionSimulator
 from zipline.utils.protocol_utils import  ndict
+from zipline.utils.env_utils import stdout_only_pipe
+
+from logbook import Logger
+from logbook import queues
 
 LOGGER = logging.getLogger('ZiplineLogger')
 
 class TradeSimulationClient(Component):
 
-    def init(self, trading_environment, sim_style):
+    def init(self, trading_environment, sim_style, log_socket):
         self.received_count         = 0
         self.prev_dt                = None
         self.event_queue            = None
@@ -28,7 +33,18 @@ class TradeSimulationClient(Component):
 
         self.event_data = ndict()
         self.perf = perf.PerformanceTracker(self.trading_environment)
+        
+        self.log_socket = log_socket
 
+        #If we have a log socket,setup context managers for capturing user logs.
+        if log_socket:
+            log = Logger("User logs")
+            self.stdout_capture = stdout_only_pipe(log, 'user algo stdout')
+
+            handler = queues.ZeroMQHandler()
+            handler.socket.connect(log_socket)
+            self.zmq_out = handler.threadbound()
+        
     @property
     def get_id(self):
         return str(zp.FINANCE_COMPONENT.TRADING_CLIENT)
@@ -41,9 +57,16 @@ class TradeSimulationClient(Component):
         self.algorithm = algorithm
         # register the trading_client's order method with the algorithm
         self.algorithm.set_order(self.order)
-        # ask the algorithm to initialize
-        self.algorithm.initialize()
 
+        # ask the algorithm to initialize, routing stdout to a zmq PUB socket.
+        if self.log_socket:
+            with self.zmq_out, self.stdout_capture:
+                self.algorithm.initialize()
+        
+        # if we don't have a log socket, initialize anyway.
+        else:
+            self.algorithm.initialize()
+            
     def open(self):
         self.result_feed = self.connect_result()
         self.perf.open(self.context)
@@ -67,7 +90,7 @@ class TradeSimulationClient(Component):
 
             # result_feed is a merge component, so unframe accordingly
             event = zp.MERGE_UNFRAME(msg)
-            self.received_count += 1
+             self.received_count += 1
             # update performance and relay the event to the algorithm
             self.process_event(event)
             if self.perf.exceeded_max_loss:
@@ -136,8 +159,14 @@ class TradeSimulationClient(Component):
         self.algorithm.set_portfolio(current_portfolio)
         data = self.get_data()
         if len(data) > 0:
-            self.algorithm.handle_data(data)
-
+            # try to run algo with log rerouting
+            if log_socket:
+                with self.zmq_out, self.stdout_capture:
+                    self.algorithm.handle_data(data)
+            # if no log socket, just run the algo normally
+            else:
+                self.algorithm.handle_data(data)
+                    
     def connect_order(self):
         return self.connect_push_socket(self.addresses['order_address'])