documented algorithm protocol, split example algorithms into their own module. added filter methods to algo and to datasources. The top level zipline is responsible for piping algo's filter to the datasource.

zipline/finance/trading.py

@@ -17,11 +17,11 @@ class TradeSimulationClient(qmsg.Component):
         self.received_count         = 0
         self.prev_dt                = None
         self.event_queue            = None
-        self.event_callbacks        = []
         self.txn_count              = 0
         self.trading_environment    = trading_environment
         self.current_dt             = trading_environment.period_start
         self.last_iteration_dur     = datetime.timedelta(seconds=0)
+        self.algorithm              = None
         
         assert self.trading_environment.frame_index != None
         self.event_frame = pandas.DataFrame(
@@ -41,15 +41,15 @@ class TradeSimulationClient(qmsg.Component):
     @property
     def get_id(self):
         return str(zp.FINANCE_COMPONENT.TRADING_CLIENT)
-    
-    def add_event_callback(self, callback):
+        
+    def set_algorithm(self, algorithm):
         """
-        :param callable callback: must be a function with the signature
-        f(event), where event is a namedict whose properties depend on the
-        upstream configuration of the zipline. It will include datasource and
-        transformations.
+        :param algorithm: must implement the algorithm protocol. See 
+        algorithm_protocol.rst.
         """
-        self.event_callbacks.append(callback)
+        self.algorithm = algorithm 
+        #register the trading_client's order method with the algorithm
+        self.algorithm.set_order(self.order)
     
     def open(self):
         self.result_feed = self.connect_result()
@@ -66,7 +66,7 @@ class TradeSimulationClient(qmsg.Component):
 
             if msg == str(zp.CONTROL_PROTOCOL.DONE):
                 qutil.LOGGER.info("Client is DONE!")
-                self.run_callbacks()
+                self.run_algorithm()
                 self.signal_order_done()
                 self.signal_done()
                 return
@@ -83,7 +83,7 @@ class TradeSimulationClient(qmsg.Component):
                 self.queue_event(event)
                 
                 if event.dt >= self.current_dt:
-                    self.run_callbacks()
+                    self.run_algorithm()
                 
                 #update time based on receipt of the order
                 self.last_iteration_dur = datetime.datetime.utcnow() - event_start
@@ -93,10 +93,9 @@ class TradeSimulationClient(qmsg.Component):
             #signal done to order source.
             self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK))
             
-    def run_callbacks(self):
+    def run_algorithm(self):
         frame = self.get_frame()
-        for cb in self.event_callbacks:
-            cb(frame)
+        self.algorithm.handle_frame(frame)
     
     def connect_order(self):
         return self.connect_push_socket(self.addresses['order_address'])

zipline/lines.py

@@ -84,7 +84,7 @@ import zipline.protocol as zp
 import zipline.finance.performance as perf
 import zipline.messaging as zmsg
 
-from zipline.test.client import TestAlgorithm
+from zipline.test.algorithms import TestAlgorithm
 from zipline.sources import SpecificEquityTrades
 from zipline.finance.trading import TransactionSimulator, OrderDataSource, \
 TradeSimulationClient
@@ -114,11 +114,9 @@ class SimulatedTrading(object):
     def __init__(self, **config):
         """
         :param config: a dict with the following required properties::
-            - algorithm: a class that follows the algorithm protocol. Must
-            have a handle_frame method that accepts a pandas.Dataframe of the 
-            current state of the simulation universe. Must have an order 
-            property which can be set equal to the order method of 
-            trading_client. (TODO: where should this protocol be documented?)
+            - algorithm: a class that follows the algorithm protocol. See
+            :py:meth:`zipline.finance.trading.TradingSimulationClient.add_algorithm`
+            for details.
             - trading_environment: an instance of
             :py:class:`zipline.trading.TradingEnvironment`
             - allocator: an instance of 
@@ -149,43 +147,30 @@ class SimulatedTrading(object):
             sockets[7],
             logging = qutil.LOGGER
         )
-
+        
+        self.started = False
         
         self.sim = config['simulator_class'](addresses)
             
         self.clients = {}
         self.trading_client = TradeSimulationClient(self.trading_environment)
-        self.clients[self.trading_client.get_id] = self.trading_client
+        self.add_client(self.trading_client)
         
         # setup all sources
         self.sources = {}
         self.order_source = OrderDataSource()
-        self.sources[self.order_source.get_id] = self.order_source
+        self.add_source(self.order_source)
         
         #setup transforms
         self.transaction_sim = TransactionSimulator()
         self.transforms = {}
-        self.transforms[self.transaction_sim.get_id] = self.transaction_sim
+        self.add_transform(self.transaction_sim)
         
-        #register all components
-        self.sim.register_components([
-            self.trading_client, 
-            self.order_source, 
-            self.transaction_sim 
-            ])
-            
         self.sim.register_controller( self.con )
         self.sim.on_done = self.shutdown()
-        self.started = False
         
-        ##################################################################
-        #TODO: the next two lines of code need refactoring from RealDiehl
-        ##################################################################
-        #wire up a callback inside the algorithm to receive frames from the
-        #trading client
-        self.trading_client.add_event_callback(self.algorithm.handle_frame)
-        #register the trading_client's order method with the algorithm
-        self.algorithm.set_order(self.trading_client.order)
+        
+        self.trading_client.set_algorithm(self.algorithm)
     
     @staticmethod
     def create_test_zipline(**config):
@@ -202,7 +187,7 @@ class SimulatedTrading(object):
             subclass of ComponentHost to hold the whole zipline. Defaults to
             :py:class:`zipline.simulator.Simulator`   
             - algorithm - optional parameter providing an algorithm. defaults
-            to :py:class:`zipline.test.client.TestAlgorithm`
+            to :py:class:`zipline.test.algorithms.TestAlgorithm`
         """
         assert isinstance(config, dict)
         
@@ -270,17 +255,27 @@ class SimulatedTrading(object):
         return zipline
         
     def add_source(self, source):
+        """
+        Adds the source to the zipline, sets the sid filter of the
+        source to the algorithm's sid filter.
+        """
         assert isinstance(source, zmsg.DataSource)
         self.check_started()    
+        source.set_filter('SID', self.algorithm.get_sid_filter)
         self.sim.register_components([source])
         self.sources[source.get_id] = source
-        
     
     def add_transform(self, transform):
         assert isinstance(transform, zmsg.BaseTransform)
         self.check_started()
         self.sim.register_components([transform])
-        self.sources[transform.get_id] = transform
+        self.transforms[transform.get_id] = transform
+        
+    def add_client(self, client):
+        assert isinstance(client, TradeSimulationClient)
+        self.check_started()
+        self.sim.register_components([client])
+        self.clients[client.get_id] = client
     
     def check_started(self):
         if self.started:

zipline/messaging.py

@@ -545,16 +545,26 @@ class DataSource(Component):
     Baseclass for data sources. Subclass and implement send_all - usually this
     means looping through all records in a store, converting to a dict, and
     calling send(map).
+    
+    Every datasource has a dict property to hold filters::
+        - key -- name of the filter, e.g. SID
+        - value -- a primitive representing the filter. e.g. a list of ints.
+        
+    Modify the datasource's filters via the set_filter(name, value)
     """
     def __init__(self, source_id):
         Component.__init__(self)
 
         self.id = source_id
         self.init()
+        self.filter = {}
 
     def init(self):
         self.cur_event = None
 
+    def set_filter(self, name, value):
+        self.filter[name] = value
+
     @property
     def get_id(self):
         return self.id

zipline/test/algorithms.py

@@ -0,0 +1,82 @@
+"""
+Algorithm Protocol
+===================
+
+For a class to be passed as a trading algorithm to the 
+:py:class:`zipline.lines.SimulatedTrading` zipline
+it must follow an implementation protocol. Examples of this algorithm protocol
+are provided below.
+
+The algorithm must expose methods::
+  - get_sid_filter: method that takes no args, and returns a list
+  of valid sids. List must have a length between 1 and 10. If None is returned
+  the filter will block all events.
+  
+  - handle_frame: method that accepts a :py:class:`pandas.Dataframe` of the 
+  current state of the simulation universe. An example frame:
++-----------------+--------------+----------------+--------------------+
+|                 | SID(133)     |  SID(134)      | SID(135)           |
++=================+==============+=====================================+
+| price	          | $10.10       | $22.50         | $13.37             |
++-----------------+--------------+----------------+--------------------+
+| volume          | 10,000       | 5,000          | 50,000             |
++-----------------+--------------+----------------+--------------------+
+| mvg_avg_30      | $9.97        | $22.61         | $13.37             |
++-----------------+--------------+----------------+--------------------+
+| dt              |	6/30/2012    | 6/30/2012      | 6/29/2012          |
++-----------------+--------------+----------------+--------------------+
+            
+The algorithm must also expose settable properties:
+  - order: property which can be set equal to the order method of 
+  trading_client. An algorithm can then place orders with a valid
+  SID and a number of shares::
+      self.order(SID(133), share_count)
+
+"""
+
+import zipline.protocol as zp
+
+class TestAlgorithm():
+    """
+    This algorithm will send a specified number of orders, to allow unit tests
+    to verify the orders sent/received, transactions created, and positions
+    at the close of a simulation.
+    """
+    
+    def __init__(self, sid, amount, order_count):
+        self.count = order_count
+        self.sid = sid
+        self.amount = amount
+        self.incr = 0
+        self.done = False
+        self.order = None
+        
+    def set_order(self, order_callable):
+        self.order = order_callable
+        
+    def handle_frame(self, frame):
+        for dt, s in frame.iteritems():     
+            data = {}
+            data.update(s)
+            event = zp.namedict(data)
+            #place an order for 100 shares of sid:133
+            if self.incr < self.count:
+                self.order(self.sid, self.amount)
+                self.incr += 1
+                
+    def get_sid_filter(self):
+        return [self.sid]
+                
+class NoopAlgorithm(object):
+    """
+    Dolce fa niente.
+    """
+        
+    def set_order(self, order_callable):
+        pass
+    
+    def handle_frame(self, frame):
+        pass
+    
+    def get_sid_filter():
+        return None

zipline/test/client.py

@@ -83,36 +83,3 @@ class TestClient(qmsg.Component):
             
     def unframe(self, msg):
         return zp.MERGE_UNFRAME(msg)
-
-
-class TestAlgorithm():
-    
-    def __init__(self, sid, amount, order_count):
-        self.count = order_count
-        self.sid = sid
-        self.amount = amount
-        self.incr = 0
-        self.done = False
-        self.order = None
-        
-    def set_order(self, order_callable):
-        self.order = order_callable
-        
-    def handle_frame(self, frame):
-        for dt, s in frame.iteritems():     
-            data = {}
-            data.update(s)
-            event = zp.namedict(data)
-            #place an order for 100 shares of sid:133
-            if self.incr < self.count:
-                self.order(self.sid, self.amount)
-                self.incr += 1
-                
-class NoopAlgorithm(object):
-        
-    def set_order(self, order_callable):
-        pass
-    
-    def handle_frame(self, frame):
-        pass
-        

zipline/test/test_finance.py

@@ -14,7 +14,7 @@ import zipline.finance.risk as risk
 import zipline.protocol as zp
 import zipline.finance.performance as perf
 
-from zipline.test.client import TestAlgorithm
+from zipline.test.algorithms import TestAlgorithm
 from zipline.sources import SpecificEquityTrades
 from zipline.finance.trading import TransactionSimulator, OrderDataSource, \
 TradeSimulationClient, TradingEnvironment
@@ -86,7 +86,7 @@ class FinanceTestCase(TestCase):
             zipline.algorithm.count, 
             "The order source should have sent as many orders as the algo."
         )
-            
+        
         transaction_sim = zipline.transforms[zp.TRANSFORM_TYPE.TRANSACTION]
         self.assertEqual(
             transaction_sim.txn_count,