# ============================================================================ class LimitReader(object): """ A reader which will not read more than specified limit """ def __init__(self, stream, limit): self.stream = stream self.limit = limit self._orig_limit = limit def _update(self, buff): length = len(buff) self.limit -= length return buff def read(self, length=None): if length is not None: length = min(length, self.limit) else: length = self.limit if length == 0: return b'' buff = self.stream.read(length) return self._update(buff) def readline(self, length=None): if length is not None: length = min(length, self.limit) else: length = self.limit if length == 0: return b'' buff = self.stream.readline(length) return self._update(buff) def close(self): self.stream.close() def tell(self): # implement our own tell return self._orig_limit - self.limit @staticmethod def wrap_stream(stream, content_length): """ If given content_length is an int > 0, wrap the stream in a LimitReader. Otherwise, return the stream unaltered """ try: content_length = int(content_length) if content_length >= 0: # optimize: if already a LimitStream, set limit to # the smaller of the two limits if isinstance(stream, LimitReader): stream.limit = min(stream.limit, content_length) else: stream = LimitReader(stream, content_length) except (ValueError, TypeError): pass return stream