diff --git a/pre_commit_images/optimizer.py b/pre_commit_images/optimizer.py index b76a9f3..d6c72e9 100644 --- a/pre_commit_images/optimizer.py +++ b/pre_commit_images/optimizer.py @@ -2,39 +2,40 @@ from collections.abc import Callable from collections.abc import Iterable from pathlib import Path -from tempfile import TemporaryFile +from tempfile import SpooledTemporaryFile from typing import IO -def _optimize_images(images: Iterable[str], optimizer_fn: Callable[[Path, IO[bytes]], None], threshold: int) -> bool: - def bytes_to_readable(file_size: int) -> str: - if file_size / (1024 * 1024) > 1: - return f"{file_size / (1024 * 1024):.2f}Mb" +def _bytes_to_readable(file_size: int) -> str: + if file_size / (1024 * 1024) > 1: + return f"{file_size / (1024 * 1024):.2f}Mb" - if file_size / (1024) > 1: - return f"{(file_size/1024):.2f}Kb" + if file_size / (1024) > 1: + return f"{(file_size/1024):.2f}Kb" - return f"{file_size}b" + return f"{file_size}b" - def optimize_single_image(source: Path, temp: IO[bytes]) -> None: - optimizer_fn(source, temp) +def _optimize_images(images: Iterable[str], optimizer: Callable[[Path, IO[bytes]], None], threshold: int) -> bool: + def keep_smaller_image(source: Path, candidate: IO[bytes]) -> None: source_size = source.stat().st_size - diff = source_size - temp.tell() + diff = source_size - candidate.tell() if diff > threshold: - temp.seek(0) - source.write_bytes(temp.read()) + candidate.seek(0) + source.write_bytes(candidate.read()) - readable_diff = bytes_to_readable(diff) - readable_size = bytes_to_readable(source_size) - print(f"Optimized {str(source)} by {readable_diff} of {readable_size} ({diff/source_size:.2%})") + readable_diff = _bytes_to_readable(diff) + readable_size = _bytes_to_readable(source_size) + print(f"Optimized {source} by {readable_diff} of {readable_size} ({diff/source_size:.2%})") ret = True for image in images: try: - with TemporaryFile() as temp: - optimize_single_image(Path(image), temp) + with SpooledTemporaryFile() as temp: + image_path = Path(image) + optimizer(image_path, temp) + keep_smaller_image(image_path, temp) except Exception as exc: # noqa: PERF203 print( f"Failed optimization for {image} ({exc})",