Source

protocol-example / ptypes.py

Full commit
from inspect import currentframe
import struct


def mktoken(name, types):
    """Create a token"""
    return Token(types['java'], name, fmt=types['ctypes'], 
                 simple=not types['java'] == "String")

def typedef(type_spec):
    """Function that creates a type function,
    which associates a name with a type specification"""
    def type_fn(name):
        """Inject the type definition type_def into 
        the calling scope with the given name"""
        frame = currentframe().f_back
        frame.f_locals[name] =  mktoken(name, type_spec)       
        del frame
    return type_fn

i64 = typedef({'ctypes': 'q', 'java': 'long'})
i16 = typedef({'ctypes': 'h', 'java': 'short'})
string = typedef({'ctypes': 'x', 'java': 'String'})

class Token(object):
    """Represents one or more tokens"""
    simple = True

    def __init__(self, javaname, name, fmt="x", simple=True):
        self.simple = simple
        self.fmt = fmt
        self.javaname = javaname
        self.name = name
        self.pattern_size = struct.calcsize(self.fmt)

    def parse(self, buf):
        if self.simple:
            values = struct.unpack("!%s" % (self.fmt), 
                                   buf.data[:self.pattern_size])
            buf.data = buf.data[self.pattern_size:]
            
            return {self.name: values[0]}
        
        else:
            
            strlen = struct.unpack("!h", buf.data[:2])[0]
            str_ = struct.unpack("!%ss" % (strlen,), buf.data[2:2+strlen])[0]
            buf.data = buf.data[2+strlen:]
            return {self.name: str_} 
    
    
    def __str__(self):
        type_ = "simple" if self.simple else "string"
        return "%s token with format %s" % (type_, self.fmt)
    
class Buffer(object):
    """Class to hold a message buffer. Each token will
    consume parts of the buffer as it parses the message"""
    def __init__(self, data):
        self.data = data  
        
    
class Message(object):
    """A message, knows how to parse data"""
    def __init__(self, msgname, msgdesc, msgtype, *tokens):
        self.tokens = tokens
        self.__doc__ = msgdesc
        self.parsers = self.tokens
        self.name = msgname
        self.type = hex(msgtype) #for java
 
        
    def parse(self, data):
        buf = Buffer(data[1:]) #remove type byte
        result = {'type': self.name}
        for res in [x.parse(buf) for x in self.parsers]:            
            result.update(res)
        assert len(buf.data) == 0
        return result