length-limited file-like object

sizeboundfile.py

from cStringIO import StringIO


class SizeBoundFile(object):
    """a length-limiting file object wrapper

    it reads from the underlying file object but never past the given size, so
    it is suitable for use around a socket.makefile() file given a particular
    Content-Length -- reading from this simulates EOF rather than reading into
    the next message.
    """
    CHUNKSIZE = 4096
    NEWLINE = '\n'

    def __init__(self, size, file_obj):
        self.size = size
        self._file_obj = file_obj
        self._rbuf = StringIO()

        # self._remaining tracks the # of bytes remaining to be read from the
        # underlying socket -- buffered bytes don't count, so this isn't
        # necessarily the # of bytes remaining to be handed out in read()s
        self._remaining = size

    def _read_some(self, nbytes):
        buffered = self._rbuf.getvalue()
        if buffered:
            self._rbuf.seek(0)
            self._rbuf.truncate()
            self._rbuf.write(buffered[nbytes:])
            return buffered[:nbytes]

        chunk = min(self.CHUNKSIZE, self._remaining)
        data = self._file_obj.read(chunk)
        if len(data) < chunk:
            # *real* EOF
            self._remaining = 0
            data

        self._remaining -= len(data)
        self._rbuf.write(data[nbytes:])
        return data[:nbytes]

    def read(self, nbytes=-1):
        if nbytes < 0:
            # read until EOF (real or simulated)
            results = StringIO()
            while 1:
                data = self._read_some(self.CHUNKSIZE)
                if not data:
                    return results.getvalue()
                results.write(data)

        return self._read_some(nbytes)

    def readline(self):
        data, index = "", -1
        buf = StringIO()

        while index < 0:
            buf.write(data)
            data = self._read_some(self.CHUNKSIZE)
            index = data.find(self.NEWLINE)

        index += len(self.NEWLINE)
        self._rbuf.write(data[index:])
        return buf.getvalue() + data[:index]

    def readlines(self):
        line = self.readline()
        while line:
            yield line
            line = self.readline()