autoytarchive/autoytarchive/cappedio.py

import io
import os

class CappedBufferedReader(io.BufferedReader):
    def __init__(self, raw, buffer_size=io.DEFAULT_BUFFER_SIZE, capped_size=None):
        super().__init__(raw, buffer_size)
        self.capped_size = capped_size

    def read(self, size=None):
        if self.capped_size is not None:
            if size is None or size < 0 or size > self.capped_size:
                size = self.capped_size
            self.capped_size -= size
        return super().read(size)

    def seek(self, offset, whence=os.SEEK_SET):
        if self.capped_size is not None:
            if offset == 0 and whence == os.SEEK_END:
                if self.capped_size < 0:
                    offset = 0
                else:
                    real_end = super().seek(0, os.SEEK_END)
                    if self.capped_size > real_end:
                        offset = real_end
                    else:
                        offset = self.capped_size
                self.capped_size -= offset
                whence = os.SEEK_SET
            elif whence == os.SEEK_SET:
                current_pos = self.tell()
                if current_pos > offset:
                    self.capped_size += current_pos - offset
        return super().seek(offset, whence)

    @property
    def name(self):
        try:
            return self._name
        except AttributeError:
            return super().name

    @name.setter
    def name(self, new_name):
        self._name = new_name

def bopen(file, capped_size):
    return CappedBufferedReader(open(file, 'rb', buffering=0), capped_size=capped_size)