bakcode/python/IntHeaderProtobufFile.py

#!/usr/bin/env python
#encoding:utf-8

import string
import struct
from google.protobuf import message as _message

class Reader(object):
    
    def __init__(self, protobufClass, file):
        self.protobufClass = protobufClass
        self.file = file
        self.f = open(self.file, "rb")

    def __enter__(self):
        pass

    def __iter__(self):
        return self
    
    def next(self):
        header = self.f.read(4)
        if len(header) < 4:
            self.__exit__(None, None, None)
            raise StopIteration()
        length = struct.unpack('>i', header)
        bytes = self.f.read(length[0])
        protoObj = self.protobufClass()
        protoObj.ParseFromString(bytes)
        return protoObj
    
    def __next__(self):
        self.next()
            
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.f.close()

class Writer(object):
    
    def __init__(self, file):
        self.file = file
        self.f = open(self.file, "wb")

    def __enter__(self):
        pass

    def write(self, message):
        if isinstance(message, str):
            length = len(message)
            self.f.write(struct.pack(">i", length))
            self.f.write(message)
            return
        if isinstance(message, _message.Message):
            bytes = message.SerializeToString()
            length = len(bytes)
            self.f.write(struct.pack(">i", length))
            self.f.write(bytes)
            return
        
    def flush(self):
        self.f.flush()

    def close(self):
        self.f.close()
    
    def isclosed(self):
        return self.f.closed

    closed = property(isclosed)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

if __name__ == "__main__":
    pass