root/trunk/storage/sockets.py

import fixedpoint
try:
    import cPickle as pickle
except ImportError:
    import pickle
import socket
import select


def stream(unit):
    data = {}
    if isinstance(unit, dict):
        for key, value in unit.iteritems():
            if isinstance(value, fixedpoint.FixedPoint):
                # Unpicklable
                data[key] = (value.n, value.p)
            else:
                data[key] = value
    else:
        # Assume it's a dejavu.Unit.
        for key in unit.__class__.properties():
            value = getattr(unit, key)
            if isinstance(value, fixedpoint.FixedPoint):
                # Unpicklable
                data[key] = (value.n, value.p)
            else:
                data[key] = value
    return pickle.dumps(data)

def destream(unit, data):
    # data will be a pickled dictionary of properties for a unit.
    attrdict = pickle.loads(data)
    cls = unit.__class__
    for key in cls.properties():
        value = attrdict[key]
        if cls.property_type(key) == fixedpoint.FixedPoint:
            if value is not None:
                n, p = value
                value = fixedpoint.FixedPoint(0, 2)
                value.n, value.p = n, p
        # Set the attribute directly to avoid __set__ overhead.
        unit._properties[key] = value

def dechunk(data):
    """Extract chunks from a stream of data."""
    chunks = []
    while data:
        size = (256 * int(ord(data[0]))) + int(ord(data[1]))
        chunks.append(data[2:size + 2])
        data = data[size + 2:]
    return chunks

def sendall(conn, msg):
    size = len(msg)
    if size > (256 * 256):
        raise ValueError("msg to send is too large (%s)." % size)
    conn.sendall(chr(size >> 8))
    conn.sendall(chr(size & 0xFF))
    conn.sendall(msg)


class SocketClient(object):
    """Client (dejavu) end of a socket for handling dejavu Units."""
    
    # For some reason, the default of '' for localhost doesn't work on Win2k.
    def __init__(self, host='127.0.0.1', port=51111, timeout=5):
        self.host = host
        self.port = port
        self.timeout = timeout
    
    def query(self, msg):
        """Send msg to the peer and return the response(s) in a list.
        
        Socket responses specify their length in the first two
        bytes (hi, lo). Then, the data follows. If the socket wishes to
        return another "object", it will follow the data with another
        length+data stream.
        """
        conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
##        conn.settimeout(self.timeout)
        conn.connect((self.host, self.port))
        sendall(conn, msg)
        conn.shutdown(1)
        
        # Receive reply.
        data = []
        while True:
            try:
                chunk = conn.recv(1024)
            except Exception, x:
                if x.args[0] != 10035:
                    raise x
            else:
                if chunk == '':
                    break
                data.append(chunk)
        
        conn.close()
        
        data = dechunk(''.join(data))
        if data and data[0] == 'ERROR':
            raise pickle.loads(data[1])
        return data


class SocketServer(object):
    """Server end of a socket for handling dejavu Units.
    
    Use this class to build wrappers for non-standard databases or
    other sources of dejavu Units.
    """
    
    def open(self, host='127.0.0.1', port=51111, backlog=5):
        """open(host='127.0.0.1', port=51111, backlog=5)."""
        self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
##        self.conn.setblocking(0)
        self.conn.bind((host, port))
        self.conn.listen(backlog)
    
    def read(self):
        """read(). Read from the socket."""
        readables, writables, e = select.select([self.conn], [self.conn], [], 1)
        if self.conn in readables:
            (sockobj, address) = self.conn.accept()
            
            data = []
            while True:
                try:
                    chunk = sockobj.recv(1024)
                except Exception, x:
                    if x.args[0] != 10035:
                        raise x
                else:
                    if chunk == '':
                        break
                    data.append(chunk)
            return sockobj, dechunk(''.join(data))[0]
        else:
            return None, ''
    
    def close(self):
        self.conn.close()