diff --git a/pywb/utils/io.py b/pywb/utils/io.py index c5de3ede..a659518d 100644 --- a/pywb/utils/io.py +++ b/pywb/utils/io.py @@ -121,3 +121,18 @@ class OffsetLimitReader(LimitReader): def readline(self, length=None): self._skip() return super(OffsetLimitReader, self).readline(length) + + +# ============================================================================ +class StreamClosingReader(object): + def __init__(self, stream): + self.stream = stream + + def read(self, length=None): + return self.stream.read(length) + + def readline(self, length=None): + return self.stream.readline(length) + + def close(self): + no_except_close(self.stream) diff --git a/pywb/utils/loaders.py b/pywb/utils/loaders.py index e2933811..06e88e38 100644 --- a/pywb/utils/loaders.py +++ b/pywb/utils/loaders.py @@ -21,7 +21,7 @@ import re from io import open, BytesIO from warcio.limitreader import LimitReader -from pywb.utils.io import no_except_close +from pywb.utils.io import no_except_close, StreamClosingReader try: import boto3 @@ -355,7 +355,7 @@ class HttpLoader(BaseLoader): r = self.session.get(url, headers=headers, stream=True) r.raise_for_status() - return r.raw + return StreamClosingReader(r.raw) # ================================================================= diff --git a/pywb/warcserver/resource/blockrecordloader.py b/pywb/warcserver/resource/blockrecordloader.py index edf98a95..a266318c 100644 --- a/pywb/warcserver/resource/blockrecordloader.py +++ b/pywb/warcserver/resource/blockrecordloader.py @@ -2,7 +2,7 @@ from warcio.bufferedreaders import DecompressingBufferedReader from warcio.recordloader import ArcWarcRecordLoader from pywb.utils.loaders import BlockLoader -from pywb.utils.io import BUFF_SIZE, no_except_close +from pywb.utils.io import BUFF_SIZE #================================================================= @@ -32,8 +32,4 @@ class BlockArcWarcRecordLoader(ArcWarcRecordLoader): decomp_type=decomp_type, block_size=self.block_size) - res = self.parse_record_stream(stream, no_record_parse=no_record_parse) - - no_except_close(stream) - - return res + return self.parse_record_stream(stream, no_record_parse=no_record_parse)