Merge pull request #48 from quantopian/change-trade-from-ndict-to-object-based-event

Changes tests from using an ndict for trades to an Event object.

zipline/finance/performance.py

@@ -252,7 +252,6 @@ class PerformanceTracker(object):
 
         message = None
 
-        assert isinstance(event, ndict)
         self.event_count += 1
 
         if(event.dt > self.market_close):

zipline/gens/sort.py

@@ -19,7 +19,6 @@ Sorting generator.
 import logbook
 
 from collections import deque
-from zipline import ndict
 from zipline.gens.utils import (
     assert_datasource_unframe_protocol,
     assert_sort_protocol
@@ -127,9 +126,6 @@ def pop_oldest(sources):
 
 # Return the event with the older timestamp.  Break ties by source_id.
 def older(oldest, current):
-    assert isinstance(oldest, ndict)
-    assert isinstance(oldest, ndict)
-
     # Try to compare by dt.
     if oldest.dt < current.dt:
         return oldest

zipline/gens/utils.py

@@ -22,7 +22,10 @@ from hashlib import md5
 from datetime import datetime
 from itertools import izip_longest
 from zipline import ndict
-from zipline.protocol import DATASOURCE_TYPE
+from zipline.protocol import (
+    DATASOURCE_TYPE,
+    Event
+)
 
 
 def mock_raw_event(sid, dt):
@@ -91,25 +94,25 @@ def hash_args(*args, **kwargs):
 
 def create_trade(sid, price, amount, datetime, source_id="test_factory"):
 
-    row = ndict({
-        'source_id': source_id,
-        'type': DATASOURCE_TYPE.TRADE,
-        'sid': sid,
-        'dt': datetime,
-        'price': price,
-        'close': price,
-        'open': price,
-        'low': price * .95,
-        'high': price * 1.05,
-        'volume': amount
-    })
-    return row
+    trade = Event()
+
+    trade.source_id = source_id
+    trade.type = DATASOURCE_TYPE.TRADE
+    trade.sid = sid
+    trade.dt = datetime
+    trade.price = price
+    trade.close = price
+    trade.open = price
+    trade.low = price * .95
+    trade.high = price * 1.05
+    trade.volume = amount
+
+    return trade
 
 
 def assert_datasource_protocol(event):
     """Assert that an event meets the protocol for datasource outputs."""
 
-    assert isinstance(event, ndict)
     assert isinstance(event.source_id, basestring)
     assert event.type in DATASOURCE_TYPE
 
@@ -123,7 +126,6 @@ def assert_trade_protocol(event):
     """Assert that an event meets the protocol for datasource TRADE outputs."""
     assert_datasource_protocol(event)
 
-    assert isinstance(event, ndict)
     assert event.type == DATASOURCE_TYPE.TRADE
     assert isinstance(event.sid, int)
     assert isinstance(event.price, numbers.Real)
@@ -133,35 +135,23 @@ def assert_trade_protocol(event):
 
 def assert_datasource_unframe_protocol(event):
     """Assert that an event is valid output of zp.DATASOURCE_UNFRAME."""
-    assert isinstance(event, ndict)
     assert isinstance(event.source_id, basestring)
     assert event.type in DATASOURCE_TYPE
-    assert 'dt' in event
 
 
 def assert_sort_protocol(event):
     """Assert that an event is valid input to zp.FEED_FRAME."""
-    assert isinstance(event, ndict)
     assert isinstance(event.source_id, basestring)
     assert event.type in DATASOURCE_TYPE
-    assert 'dt' in event
 
 
 def assert_sort_unframe_protocol(event):
     """Same as above."""
-    assert isinstance(event, ndict)
     assert isinstance(event.source_id, basestring)
     assert event.type in DATASOURCE_TYPE
-    assert 'dt' in event
-
-
-def assert_transform_protocol(event):
-    """Transforms should return an ndict to be merged by merge."""
-    assert isinstance(event, ndict)
 
 
 def assert_merge_protocol(tnfm_ids, message):
     """Merge should output an ndict with a field for each id
     in its transform set."""
-    assert isinstance(message, ndict)
     assert set(tnfm_ids) == set(message.keys())

zipline/protocol.py

@@ -27,3 +27,21 @@ DATASOURCE_TYPE = Enum(
     'EMPTY',
     'DONE'
 )
+
+
+class Event(object):
+
+    def __getitem__(self, name):
+        return getattr(self, name)
+
+    def __setitem__(self, name, value):
+        setattr(self, name, value)
+
+    def __delitem__(self, name):
+        delattr(self, name)
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__

zipline/transforms/mavg.py

@@ -139,8 +139,6 @@ class MovingAverageEventWindow(EventWindow):
         We only allow events with all of our tracked fields.
         """
         for field in self.fields:
-            assert field in event, \
-                "Event missing [%s] in MovingAverageEventWindow" % field
             assert isinstance(event[field], Number), \
                 "Got %s for %s in MovingAverageEventWindow" % (event[field],
                                                                field)

zipline/transforms/returns.py

@@ -32,8 +32,6 @@ class Returns(object):
         """
         Update and return the calculated returns for this event's sid.
         """
-        assert 'dt' in event
-        assert 'price' in event
         tracker = self.mapping[event.sid]
         tracker.update(event)
 

zipline/transforms/stddev.py

@@ -88,14 +88,12 @@ class MovingStandardDevWindow(EventWindow):
         self.sum_sqr = 0.0
 
     def handle_add(self, event):
-        assert 'price' in event
         assert isinstance(event.price, Number)
 
         self.sum += event.price
         self.sum_sqr += event.price ** 2
 
     def handle_remove(self, event):
-        assert 'price' in event
         assert isinstance(event.price, Number)
 
         self.sum -= event.price

zipline/transforms/utils.py

@@ -291,8 +291,6 @@ class EventWindow(object):
     # All event windows expect to receive events with datetime fields
     # that arrive in sorted order.
     def assert_well_formed(self, event):
-        assert isinstance(event, ndict), "Bad event in EventWindow:%s" % event
-        assert 'dt' in event, "Missing dt in EventWindow:%s" % event
         assert isinstance(event.dt, datetime), \
             "Bad dt in EventWindow:%s" % event
         if len(self.ticks) > 0: