mirror of
https://github.com/webrecorder/pywb.git
synced 2025-03-15 08:04:49 +01:00
refactor: split LimitReader into limitreader.py
This commit is contained in:
parent
b7285b1a77
commit
2b3fde028f
69
pywb/utils/limitreader.py
Normal file
69
pywb/utils/limitreader.py
Normal file
@ -0,0 +1,69 @@
|
||||
# ============================================================================
|
||||
class LimitReader(object):
|
||||
"""
|
||||
A reader which will not read more than specified limit
|
||||
"""
|
||||
|
||||
def __init__(self, stream, limit):
|
||||
self.stream = stream
|
||||
self.limit = limit
|
||||
|
||||
if hasattr(stream, 'tell'):
|
||||
self.tell = self._tell
|
||||
|
||||
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):
|
||||
return self.stream.tell()
|
||||
|
||||
@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
|
||||
|
@ -17,6 +17,7 @@ import base64
|
||||
import cgi
|
||||
|
||||
from io import open, BytesIO
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
|
||||
try:
|
||||
from boto import connect_s3
|
||||
@ -500,78 +501,6 @@ class HMACCookieMaker(object):
|
||||
return cookie
|
||||
|
||||
|
||||
#=================================================================
|
||||
# Limit Reader
|
||||
#=================================================================
|
||||
class LimitReader(object):
|
||||
"""
|
||||
A reader which will not read more than specified limit
|
||||
"""
|
||||
|
||||
def __init__(self, stream, limit):
|
||||
self.stream = stream
|
||||
self.limit = limit
|
||||
|
||||
if hasattr(stream, 'tell'):
|
||||
self.tell = self._tell
|
||||
|
||||
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):
|
||||
return self.stream.tell()
|
||||
|
||||
@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
|
||||
|
||||
|
||||
# ============================================================================
|
||||
BlockLoader.init_default_loaders()
|
||||
|
||||
|
@ -101,7 +101,7 @@ Zero-Length chunk:
|
||||
from io import BytesIO
|
||||
from pywb.utils.bufferedreaders import ChunkedDataReader, ChunkedDataException
|
||||
from pywb.utils.bufferedreaders import DecompressingBufferedReader
|
||||
from pywb.utils.loaders import LimitReader
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
|
||||
from pywb import get_test_dir
|
||||
|
||||
|
@ -141,8 +141,9 @@ from io import BytesIO
|
||||
import requests
|
||||
|
||||
from pywb.utils.loaders import BlockLoader, HMACCookieMaker, to_file_url
|
||||
from pywb.utils.loaders import LimitReader, extract_client_cookie, extract_post_query
|
||||
from pywb.utils.loaders import extract_client_cookie, extract_post_query
|
||||
from pywb.utils.loaders import append_post_query, read_last_line
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
|
||||
from pywb.utils.bufferedreaders import DecompressingBufferedReader
|
||||
|
||||
|
@ -4,7 +4,7 @@ from pywb.utils.statusandheaders import StatusAndHeaders
|
||||
from pywb.utils.statusandheaders import StatusAndHeadersParser
|
||||
from pywb.utils.statusandheaders import StatusAndHeadersParserException
|
||||
|
||||
from pywb.utils.loaders import LimitReader
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
from pywb.utils.loaders import to_native_str
|
||||
|
||||
from pywb.utils.wbexception import WbException
|
||||
|
@ -1,5 +1,5 @@
|
||||
from pywb.utils.loaders import extract_post_query, append_post_query
|
||||
from pywb.utils.loaders import LimitReader
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
from pywb.utils.statusandheaders import StatusAndHeadersParser
|
||||
|
||||
from six.moves.urllib.parse import urlsplit, quote
|
||||
|
@ -1,5 +1,5 @@
|
||||
from pywb.utils.statusandheaders import StatusAndHeaders
|
||||
from pywb.utils.loaders import LimitReader
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
from pywb.framework.cache import create_cache
|
||||
|
||||
from tempfile import NamedTemporaryFile, mkdtemp
|
||||
|
@ -7,7 +7,7 @@ from itertools import chain
|
||||
|
||||
from pywb.utils.statusandheaders import StatusAndHeaders
|
||||
from pywb.utils.wbexception import WbException, NotFoundException
|
||||
from pywb.utils.loaders import LimitReader
|
||||
from pywb.utils.limitreader import LimitReader
|
||||
from pywb.utils.timeutils import timestamp_now
|
||||
|
||||
from pywb.framework.wbrequestresponse import WbResponse
|
||||
|
Loading…
x
Reference in New Issue
Block a user