diff --git a/src/servestatic/compress.py b/src/servestatic/compress.py index 08589a8..8458d02 100644 --- a/src/servestatic/compress.py +++ b/src/servestatic/compress.py @@ -1,10 +1,10 @@ from __future__ import annotations import argparse -import concurrent.futures import gzip import os import re +from concurrent.futures import ThreadPoolExecutor, as_completed from io import BytesIO try: @@ -123,12 +123,6 @@ def write_data(path, data, suffix, stat_result): os.utime(filename, (stat_result.st_atime, stat_result.st_mtime)) return filename - def files_to_compress(self, root): - for dirpath, _dirs, files in os.walk(root): - for filename in files: - if self.should_compress(filename): - yield os.path.join(dirpath, filename) - def main(argv=None): parser = argparse.ArgumentParser( @@ -167,8 +161,17 @@ def main(argv=None): quiet=args.quiet, ) - with concurrent.futures.ThreadPoolExecutor() as executor: - executor.map(compressor.compress, compressor.files_to_compress(args.root)) + futures = [] + with ThreadPoolExecutor() as executor: + for dirpath, _dirs, files in os.walk(args.root): + futures.extend( + executor.submit(compressor.compress, os.path.join(dirpath, filename)) + for filename in files + if compressor.should_compress(filename) + ) + # Trigger any errors + for future in as_completed(futures): + future.result() return 0 diff --git a/tests/test_compress.py b/tests/test_compress.py index cf76c69..77ab845 100644 --- a/tests/test_compress.py +++ b/tests/test_compress.py @@ -6,6 +6,7 @@ import re import shutil import tempfile +from unittest import mock import pytest @@ -78,3 +79,13 @@ def test_compress(): def test_compressed_effectively_no_orig_size(): compressor = Compressor(quiet=True) assert not compressor.is_compressed_effectively("test_encoding", "test_path", 0, "test_data") + + +def test_main_error(files_dir): + with ( + pytest.raises(ValueError) as excinfo, + mock.patch.object(Compressor, "compress", side_effect=ValueError("woops")), + ): + compress_main([files_dir, "--quiet"]) + + assert excinfo.value.args == ("woops",)