You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
autoytarchive/autoytarchive/cappedio.py

47 lines
1.6 KiB

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
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)