catalyst/backtest/util.py

"""
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))