reorganized packages, test passing again

simulator/backtest/util.py

@@ -0,0 +1,162 @@
+"""
+Small classes to assist with db access, timezone calculations, and so on.
+"""
+
+import datetime
+import pytz
+import json
+import logging
+import uuid
+import zmq
+
+class DocWrap():
+    """
+        Provides attribute access style on top of dictionary results from pymongo. 
+        Allows you to access result['field'] as result.field.
+        Aliases result['_id'] to result.id.
+        
+    """
+    def __init__(self, store=None):
+        if(store == None):
+            self.store = {}
+        else:
+            self.store = store.copy()
+        if(self.store.has_key('_id')):
+            self.store['id'] = self.store['_id']
+            del(self.store['_id'])
+        
+    def __setitem__(self,key,value):
+        if(key == '_id'):
+            self.store['id'] = value
+        else:
+            self.store[key] = value
+        
+    def __getitem__(self, key):
+        if self.store.has_key(key):
+            return self.store[key]
+            
+    def __getattr__(self,attrname):
+        if self.store.has_key(attrname):
+            return self.store[attrname]
+        else:
+            raise AttributeError("No attribute named {name}".format(name=attrname))
+            
+def parse_date(dt_str):
+    """parse strings according to the same format as generated by format_date"""
+    if(dt_str == None):
+        return None
+    parts = dt_str.split(".")
+    dt = datetime.datetime.strptime(parts[0], '%Y/%m/%d-%H:%M:%S').replace(microsecond=int(parts[1]+"000")).replace(tzinfo = pytz.utc)
+    return dt
+    
+def format_date(dt):
+    """Format the date into a date with millesecond resolution and string/alphabetical sorting that is equivalent to datetime sorting"""
+    if(dt == None):
+        return None
+    dt_str = dt.strftime('%Y/%m/%d-%H:%M:%S') + "." + str(dt.microsecond / 1000)
+    return dt_str
+    
+    
+class ParallelBuffer(object):
+    """ holds several queues of events by key, allows retrieval in date order or by merging"""
+    def __init__(self, key_list):
+        self.out_socket = None
+        self.sent_count = 0
+        self.received_count = 0
+        self.draining = False
+        self.data_buffer = {}
+        for key in key_list:
+            self.data_buffer[key] = []
+    
+    def __len__(self):
+        return len(self.data_buffer)        
+    
+    def append(self, key, value):
+        self.data_buffer[key].append(value)
+        self.received_count += 1
+    
+    def next(self):
+        if(not(self.is_full() or self.draining)):
+            return
+            
+        cur = None
+        earliest = None
+        for source, events in self.data_buffer.iteritems():
+            if len(events) == 0:
+                continue
+            cur = events
+            if(earliest == None) or (cur[0]['dt'] <= earliest[0]['dt']):
+                earliest = cur
+        
+        if(earliest != None):
+            return earliest.pop(0)
+        
+    def is_full(self):
+        for source, events in self.data_buffer.iteritems():
+            if (len(events) == 0):
+                return False
+        return True
+    
+    def pending_messages(self):
+        total = 0
+        for source, events in self.data_buffer.iteritems():
+            total += len(events)
+        return total       
+        
+    def drain(self):
+        self.draining = True
+        while(self.pending_messages() > 0):
+            self.send_next()
+            
+    def send_next(self):
+        if(not(self.is_full() or self.draining)):
+            return
+            
+        event = self.next()
+        if(event != None):        
+            self.out_socket.send(json.dumps(event))
+            self.sent_count += 1   
+    
+    
+class MergedParallelBuffer(ParallelBuffer):
+    
+    def __init__(self, keys):
+        ParallelBuffer.__init__(self, keys)
+        self.feed = []
+        self.data_buffer["feed"] = self.feed
+    
+    def next(self):
+        if(not(self.is_full() or self.draining)):
+            return
+        
+        result = self.feed.pop(0)
+        for source, events in self.data_buffer.iteritems():
+            if(source == "feed"):
+                continue
+            if(len(events) > 0):
+                cur = events.pop(0)
+                result[source] = cur['value']
+        return result
+        
+        
+class FeedSync(object):
+    
+    def __init__(self, feed, name):
+        self.feed = feed
+        self.id = "{name}-{id}".format(name=name, id=uuid.uuid1())
+        self.feed.register_sync(self.id)
+        self.logger = logging.getLogger()
+        #self.logger.info("registered {id} with feed".format(id=self.id))
+        
+    def confirm(self):
+        context = zmq.Context()
+        #synchronize with feed
+        sync_socket = context.socket(zmq.REQ)
+        sync_socket.connect(self.feed.sync_address)
+        # send a synchronization request to the feed
+        sync_socket.send(self.id)
+        # wait for synchronization reply from the feed
+        sync_socket.recv()
+        sync_socket.close()
+        context.term()
+        self.logger.info("sync'd feed from {id}".format(id = self.id))