1
0
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:
Ilya Kreymer 2017-03-01 15:13:32 -08:00
parent b7285b1a77
commit 2b3fde028f
8 changed files with 77 additions and 78 deletions

69
pywb/utils/limitreader.py Normal file
View 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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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