Stream Recompressor

test/test_recompressor.py

@@ -0,0 +1,91 @@
+from . import get_test_file
+from warcio.recompressor import Recompressor, StreamRecompressor
+
+import gzip
+import pytest
+
+def test_recompress_chunked(capsys,tmp_path):
+    test_file = get_test_file('example-resource.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    recompressor = Recompressor(test_file, str(tmp_file), verbose=True)
+    recompressor.recompress()
+    out, err = capsys.readouterr()
+    assert len(out) > 0
+    assert "Records successfully read and compressed" in out
+    assert "3 records read and recompressed to file" in out
+    assert "No Errors Found!" in out
+
+def test_recompress_non_chunked(capsys,tmp_path):
+    test_file = get_test_file('example-bad-non-chunked.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    recompressor = Recompressor(test_file, str(tmp_file), verbose=True)
+    recompressor.recompress()
+    out, err = capsys.readouterr()
+    assert len(out) > 0
+    assert "ERROR: non-chunked gzip file detected" in out
+    assert "Records successfully read and compressed" in out
+    assert "6 records read and recompressed to file" in out
+    assert "Compression Errors Found and Fixed!" in out
+
+def test_recompress_chunked_decompressed_stream(tmp_path):
+    """Open a stream with befor feeding it to the _load_and_write_stream method as stream."""
+    test_file = get_test_file('example-resource.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    recompressor = Recompressor(test_file, str(tmp_file), verbose=True)
+    with open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        count = recompressor._load_and_write_stream(input, output)
+    assert count == 3
+
+def test_recompress_non_chunked_decompressed_stream_fails(tmp_path):
+    """Open a stream of a non chunked gzip befor feeding it to the _load_and_write_stream method as stream. Expect it to fail."""
+    test_file = get_test_file('example-bad-non-chunked.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    recompressor = Recompressor(test_file, str(tmp_file), verbose=True)
+    with open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        with pytest.raises(Exception):
+            recompressor._load_and_write_stream(input, output)
+
+def test_recompress_non_chunked_decompressed_stream(tmp_path):
+    """Uncompress a stream with gzip befor feeding it to the _load_and_write_stream method as stream."""
+    test_file = get_test_file('example-bad-non-chunked.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    recompressor = Recompressor(test_file, str(tmp_file), verbose=True)
+    with gzip.open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        count = recompressor._load_and_write_stream(input, output)
+    assert count == 6
+
+def test_stream_recompress_non_chunked_decompressed_stream(tmp_path):
+    """Uncompress a badly chunked stream with gzip befor feeding it to the StreamRecompressor's recompress method as stream."""
+    test_file = get_test_file('example-bad-non-chunked.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    with gzip.open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        recompressor = StreamRecompressor(input, output, verbose=True)
+        count = recompressor.recompress()
+    assert count == 6
+
+def test_stream_recompress_chunked_compressed_stream(tmp_path):
+    """Open a chunked compressed stream, feeding it to the StreamRecompressor's recompress method as stream."""
+    test_file = get_test_file('example-resource.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    with open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        recompressor = StreamRecompressor(input, output, verbose=True)
+        count = recompressor.recompress()
+    assert count == 3
+
+def test_stream_decompress_recompress_non_chunked_compressed_stream(tmp_path):
+    """Open a badly chunked stream and feed it to the StreamRecompressor's decompress_recompress method as stream."""
+    test_file = get_test_file('example-bad-non-chunked.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    with open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        recompressor = StreamRecompressor(input, output, verbose=True)
+        count = recompressor.decompress_recompress()
+    assert count == 6
+
+def test_stream_decompress_recompress_chunked_stream(tmp_path):
+    """Open a chunked compressed stream and feed it to the StreamRecompressor's decompress_recompress method as stream."""
+    test_file = get_test_file('example-resource.warc.gz')
+    tmp_file = tmp_path / "output.warc.gz"
+    with open(test_file, "rb") as input, open(tmp_file, "wb") as output:
+        recompressor = StreamRecompressor(input, output, verbose=True)
+        count = recompressor.decompress_recompress()
+    assert count == 3

warcio/recompressor.py

@@ -8,10 +8,12 @@
 import shutil
 import traceback
 import sys
+import gzip
 
 
 # ============================================================================
 class Recompressor(object):
+    """The Recompressor attempts to read a .warc or .warc.gz file and writes it to a propperly compressed .warc.gz file."""
     def __init__(self, filename, output, verbose=False):
         self.filename = filename
         self.output = output
@@ -24,16 +26,18 @@ def recompress(self):
             msg = ''
             with open(self.filename, 'rb') as stream:
                 try:
-                    count = self.load_and_write(stream, self.output)
-                    msg = 'No Errors Found!'
+                    with open(self.output, 'wb') as out:
+                        count = self._load_and_write_stream(stream, out)
+                        msg = 'No Errors Found!'
                 except Exception as e:
                     if self.verbose:
                         print('Parsing Error(s) Found:')
                         print(str(e) if isinstance(e, ArchiveLoadFailed) else repr(e))
                         print()
 
-                    count = self.decompress_and_recompress(stream, self.output)
-                    msg = 'Compression Errors Found and Fixed!'
+                    with open(self.output, 'wb') as out, tempfile.TemporaryFile() as tout:
+                        count = self._decompress_and_recompress_stream(stream, out, tout)
+                        msg = 'Compression Errors Found and Fixed!'
 
                 if self.verbose:
                     print('Records successfully read and compressed:')
@@ -53,31 +57,54 @@ def recompress(self):
             sys.exit(1)
 
     def load_and_write(self, stream, output):
-        count = 0
+        """Iterate the WARC stream to load it and write it to the output file."""
         with open(output, 'wb') as out:
-            writer = WARCWriter(filebuf=out, gzip=True)
+            return self._load_and_write_stream(stream, out)
 
-            for record in ArchiveIterator(stream,
-                                          no_record_parse=False,
-                                          arc2warc=True,
-                                          verify_http=False):
+    def decompress_and_recompress(self, stream, output):
+        """Decompress the WARC stream to a temporary location, load it, and write the recompressed result to the output file."""
+        with open(output, 'wb') as out, tempfile.TemporaryFile() as tout:
+            return self._decompress_and_recompress_stream(stream, out, tout)
 
-                writer.write_record(record)
-                count += 1
+    def _load_and_write_stream(self, in_stream, out_stream):
+        """Iterate the WARC stream to load it and write it as compressed chunked .warc.gz steam to the output stream."""
+        count = 0
+        writer = WARCWriter(filebuf=out_stream, gzip=True)
 
-            return count
+        for record in ArchiveIterator(in_stream,
+                                      no_record_parse=False,
+                                      arc2warc=True,
+                                      verify_http=False):
 
-    def decompress_and_recompress(self, stream, output):
-        with tempfile.TemporaryFile() as tout:
-            decomp = DecompressingBufferedReader(stream, read_all_members=True)
+            writer.write_record(record)
+            count += 1
+
+        return count
 
-            # decompress entire file to temp file
-            stream.seek(0)
-            shutil.copyfileobj(decomp, tout)
+    def _decompress_and_recompress_stream(self, in_stream, out_stream, tmp_stream):
+        """Decompress a WARC stream (in_stream, must be seekable) to a temporary location (tmp_stream, must be seekable), load it, and write the recompressed result to the output file."""
+        decomp = DecompressingBufferedReader(in_stream, read_all_members=True)
 
-            # attempt to compress and write temp
-            tout.seek(0)
-            return self.load_and_write(tout, output)
+        # decompress entire file to temp file
+        in_stream.seek(0)
+        shutil.copyfileobj(decomp, tmp_stream)
 
+        # attempt to compress and write temp
+        tmp_stream.seek(0)
+        return self._load_and_write_stream(tmp_stream, out_stream)
 
+class StreamRecompressor(Recompressor):
+    """The StreamRecompressor opperates on steam and attempts to read a .warc or .warc.gz stream and writes it as a propperly compressed .warc.gz stream."""
+    def __init__(self, input, output, verbose=False):
+        self.input = input
+        self.output = output
+        self.verbose = verbose
+
+    def recompress(self):
+        """Read a .warc or propperly chunked .warc.gz stream and recompresses it to proppery chunked .warc.gz stream."""
+        return self._load_and_write_stream(self.input, self.output)
 
+    def decompress_recompress(self):
+        """Reads a gzip compressed .warc stream (not necessarily propperly chunked) and recompresses it to proppery chunked .warc.gz stream."""
+        with gzip.open(self.input, "rb") as input_stream:
+            return self._load_and_write_stream(input_stream, self.output)