refactoring performance to be a plain class, not a component.

zipline/finance/performance.py

@@ -9,10 +9,9 @@ import zipline.util as qutil
 import zipline.protocol as zp
 import zipline.finance.risk as risk
 
-class PortfolioClient(qmsg.Component):
+class PerformanceTracker():
     
     def __init__(self, period_start, period_end, capital_base, trading_environment):
-        qmsg.Component.__init__(self)
         self.trading_day            = datetime.timedelta(hours=6, minutes=30)
         self.calendar_day           = datetime.timedelta(hours=24)
         self.period_start           = period_start
@@ -27,35 +26,33 @@ class PortfolioClient(qmsg.Component):
         self.capital_base           = capital_base
         self.trading_environment    = trading_environment
         self.returns                = []
-        self.cumulative_performance = PerformancePeriod(self.period_start, self.period_end, {}, 0, capital_base = capital_base)
-        self.todays_performance     = PerformancePeriod(self.market_open, self.market_close, {}, 0, capital_base = capital_base)
-    
-    @property
-    def get_id(self):
-        return str(zp.FINANCE_COMPONENT.PORTFOLIO_CLIENT)
-    
-    def open(self):
-        self.result_feed = self.connect_result()
-    
-    def do_work(self):
-        #next feed event
-        socks = dict(self.poll.poll(self.heartbeat_timeout))
-
-        if self.result_feed in socks and socks[self.result_feed] == self.zmq.POLLIN:   
-            msg = self.result_feed.recv()
-
-            if msg == str(zp.CONTROL_PROTOCOL.DONE):
-                self.handle_simulation_end()
-                qutil.LOGGER.info("Portfolio Client is DONE!")
-                self.signal_done()
-                return
-            
-            event = zp.MERGE_UNFRAME(msg)
+        self.txn_count              = 0 
+        self.event_count            = 0
+        self.cumulative_performance = PerformancePeriod(
+            self.period_start, 
+            self.period_end, 
+            {}, 
+            capital_base, 
+            capital_base = capital_base
+        )
             
+        self.todays_performance     = PerformancePeriod(
+            self.market_open, 
+            self.market_close, 
+            {}, 
+            capital_base, 
+            capital_base = capital_base
+        )
+        
+        
+    
+    def update(self, event):
+            self.event_count += 1
             if(event.dt >= self.market_close):
                 self.handle_market_close()
             
-            if event.TRANSACTION:
+            if event.TRANSACTION != None:                
+                self.txn_count += 1
                 self.cumulative_performance.execute_transaction(event.TRANSACTION)
                 self.todays_performance.execute_transaction(event.TRANSACTION)
                 
@@ -73,9 +70,7 @@ class PortfolioClient(qmsg.Component):
             #calculate performance as of last trade
             self.cumulative_performance.calculate_performance()
             self.todays_performance.calculate_performance()
-            
-            
-    
+               
     def handle_market_close(self):
         self.market_open = self.market_open + self.calendar_day
         while not self.trading_environment.is_trading_day(self.market_open):

zipline/finance/trading.py

@@ -15,11 +15,19 @@ class TradeSimulationClient(qmsg.Component):
         self.received_count     = 0
         self.prev_dt            = None
         self.event_queue        = []
+        self.event_callbacks    = []
     
     @property
     def get_id(self):
         return str(zp.FINANCE_COMPONENT.TRADING_CLIENT)
     
+    def add_event_callback(self, callback):
+        """
+        :param callable callback: must be a function with the signature
+        f(frame).
+        """
+        self.event_callbacks.append(callback)
+    
     def open(self):
         self.result_feed = self.connect_result()
         self.order_socket = self.connect_order()
@@ -39,19 +47,15 @@ class TradeSimulationClient(qmsg.Component):
                 return
             
             event = zp.MERGE_UNFRAME(msg)
-            self._handle_event(event)
+            for cb in self.event_callbacks:
+                cb(event)
+                
+            #signal done to order source.
+            self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK))
     
     def connect_order(self):
         return self.connect_push_socket(self.addresses['order_address'])
     
-    def _handle_event(self, event):
-        self.handle_event(event)
-        #signal done to order source.
-        self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK))
-    
-    def handle_event(self, event):
-        raise NotImplementedError    
-    
     def order(self, sid, amount):
         self.order_socket.send(zp.ORDER_FRAME(sid, amount))
         
@@ -151,6 +155,7 @@ class TransactionSimulator(qmsg.BaseTransform):
         qmsg.BaseTransform.__init__(self, zp.TRANSFORM_TYPE.TRANSACTION)
         self.open_orders                = {}
         self.order_count                = 0
+        self.txn_count                  = 0
         self.trade_windwo                = datetime.timedelta(seconds=30)
         self.orderTTL                   = datetime.timedelta(days=1)
         self.volume_share               = 0.05
@@ -231,7 +236,8 @@ class TransactionSimulator(qmsg.BaseTransform):
         return self.create_transaction(event.sid, amount, event.price + impact, dt.replace(tzinfo = pytz.utc), direction)
     
         
-    def create_transaction(self, sid, amount, price, dt, direction):                
+    def create_transaction(self, sid, amount, price, dt, direction):   
+        self.txn_count += 1             
         txn = {'sid'                : sid, 
                 'amount'             : int(amount), 
                 'dt'                 : dt, 

zipline/test/client.py

@@ -66,10 +66,11 @@ class TestClient(qmsg.Component):
         return zp.MERGE_UNFRAME(msg)
 
 
-class TestTradingClient(TradeSimulationClient):
+class TestAlgorithm():
     
-    def __init__(self, sid, amount, order_count):
-        TradeSimulationClient.__init__(self)
+    def __init__(self, sid, amount, order_count, trading_client):
+        self.trading_client = trading_client
+        self.trading_client.add_event_callback(self.handle_event)
         self.count = order_count
         self.sid = sid
         self.amount = amount
@@ -78,8 +79,8 @@ class TestTradingClient(TradeSimulationClient):
     def handle_event(self, event):
         #place an order for 100 shares of sid:133
         if(self.incr < self.count):
-            self.order(self.sid, self.amount)
+            self.trading_client.order(self.sid, self.amount)
             self.incr += 1
         else:
-            self.signal_order_done()
-            self.signal_done()
+            self.trading_client.signal_order_done()
+            self.trading_client.signal_done()

zipline/test/test_finance.py

@@ -10,9 +10,10 @@ import zipline.finance.risk as risk
 import zipline.protocol as zp
 import zipline.finance.performance as perf
 
-from zipline.test.client import TestTradingClient
+from zipline.test.client import TestAlgorithm
 from zipline.sources import SpecificEquityTrades
-from zipline.finance.trading import TransactionSimulator, OrderDataSource
+from zipline.finance.trading import TransactionSimulator, OrderDataSource, \
+TradeSimulationClient
 from zipline.simulator import AddressAllocator, Simulator
 from zipline.monitor import Controller
 
@@ -172,15 +173,21 @@ class FinanceTestCase(TestCase):
         )
 
         set1 = SpecificEquityTrades("flat-133", trade_history)
-
-        #client sill send 10 orders for 100 shares of 133
-        client = TestTradingClient(133, 100, 10)
+        
+        trading_client = TradeSimulationClient()
+        #client will send 10 orders for 100 shares of 133
+        test_algo = TestAlgorithm(133, 100, 10, trading_client)
         ts = datetime.strptime("02/1/2012","%m/%d/%Y").replace(tzinfo=pytz.utc)
 
         order_source = OrderDataSource(ts)
         transaction_sim = TransactionSimulator()
 
-        sim.register_components([client, order_source, transaction_sim, set1])
+        sim.register_components([
+            trading_client, 
+            order_source, 
+            transaction_sim, 
+            set1
+        ])
         sim.register_controller( con )
 
         # Simulation
@@ -242,24 +249,27 @@ class FinanceTestCase(TestCase):
         set1 = SpecificEquityTrades("flat-133", trade_history)
 
         #client sill send 10 orders for 100 shares of 133
-        client = TestTradingClient(133, 100, 10)
+        trading_client = TradeSimulationClient()
+        test_algo = TestAlgorithm(133, 100, 10, trading_client)
         ts = datetime.strptime("02/1/2012","%m/%d/%Y")
         ts = ts.replace(tzinfo=pytz.utc)
 
         order_source = OrderDataSource(ts)
         transaction_sim = TransactionSimulator()
-        portfolio_client = perf.PortfolioClient(
+        perf_tracker = perf.PerformanceTracker(
             trade_history[0]['dt'], 
             trade_history[-1]['dt'], 
             1000000.0, 
             self.trading_environment)
+        
+        #register perf_tracker to receive callbacks from the client.
+        trading_client.add_event_callback(perf_tracker.update)
     
         sim.register_components([
-            client, 
+            trading_client, 
             order_source, 
             transaction_sim, 
             set1, 
-            portfolio_client,
             ])
         sim.register_controller( con )
 
@@ -268,8 +278,34 @@ class FinanceTestCase(TestCase):
         sim_context = sim.simulate()
         sim_context.join()
 
+        
 
         # TODO: Make more assertions about the final state of the components.
         self.assertEqual(sim.feed.pending_messages(), 0, \
             "The feed should be drained of all messages, found {n} remaining." \
-            .format(n=sim.feed.pending_messages()))
+            .format(n=sim.feed.pending_messages()))
+            
+        self.assertEqual(
+            order_source.sent_count, 
+            test_algo.count, 
+            "The order source should have sent as many orders as the algo."
+        )
+            
+        self.assertEqual(
+            transaction_sim.txn_count,
+            perf_tracker.txn_count,
+            "The perf tracker should handle the same number of transactions as\
+            as the simulator emits."
+        ) 
+        
+        self.assertEqual(
+            len(perf_tracker.cumulative_performance.positions), 
+            1, 
+            "Portfolio should have one position."
+        )
+        
+        self.assertEqual(
+            perf_tracker.cumulative_performance.positions[133].sid, 
+            133, 
+            "Portfolio should have one position in 133."
+        )