Coverage for shrinky/__init__.py: 96%
141 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-23 10:01 +1000
1"""shrinky module"""
3import os
4from pathlib import Path
5import sys
6from typing import Any, Dict, Optional, Tuple, Union
8import click
9from loguru import logger
10from PIL import Image, UnidentifiedImageError
11from pillow_heif import register_heif_opener # type: ignore[import-untyped]
13register_heif_opener()
15DEFAULT_GEOMETRY = 2000
17VALID_OUTPUT_TYPES = [
18 "jpg",
19 "png",
20 "gif",
21 "webp",
22 "avif",
23 "heic",
24 "heif",
25]
28def parse_geometry(geometry_input: str) -> Tuple[Optional[int], Optional[int]]:
29 """parse the geometry provided"""
30 if "x" in geometry_input:
31 x_value, y_value = geometry_input.split("x")
32 if x_value == "":
33 x_result = None
34 else:
35 x_result = int(x_value)
37 if y_value == "":
38 y_result = None
39 else:
40 y_result = int(y_value)
41 return (x_result, y_result)
42 return (None, None)
45def set_geometry(geometry_value: Optional[str]) -> Tuple[int, int]:
46 """geometry handler"""
48 if geometry_value is None:
49 return (DEFAULT_GEOMETRY, DEFAULT_GEOMETRY)
51 max_x, max_y = parse_geometry(geometry_value)
52 if max_x is None or max_x == 0:
53 max_x = DEFAULT_GEOMETRY
54 if max_y is None or max_y == 0:
55 max_y = DEFAULT_GEOMETRY
56 logger.debug("Setting geometry to {}x{}", max_x, max_y)
57 return (max_x, max_y)
60class InvalidOutputType(Exception):
61 """raised when the output type is invalid"""
63 def __init__(self, output_type: str) -> None:
64 super().__init__(
65 f"Invalid output type: {output_type}, valid types are: {VALID_OUTPUT_TYPES}"
66 )
67 self.output_type = output_type
70def new_filename(original_filename: Path, output_type: Optional[str]) -> Path:
71 """generates a new filename based on the path"""
72 logger.debug(f"{original_filename=}")
74 basename = ".".join(original_filename.resolve().name.split(".")[:-1])
75 if output_type is not None:
76 logger.debug("Setting output type to {}", output_type.lower())
77 newname = f"{basename}.{output_type.lower()}"
78 else:
79 try:
80 extension = get_extension(original_filename)
81 if extension.lower() not in VALID_OUTPUT_TYPES:
82 raise InvalidOutputType(extension)
83 newname = f"{basename}-shrink.{extension}"
84 except ValueError:
85 newname = f"{basename}-shrink.jpg"
86 logger.error(
87 "Can't get extension for {}, setting to {}", original_filename, newname
88 )
90 return Path(f"{original_filename.parent}/{newname}").resolve()
93def get_extension(filename: Path) -> str:
94 """gets the file extension is a hacky way"""
95 if "." not in filename.name:
96 raise ValueError("Can't have an extension when there's no dot!")
97 return filename.resolve().name.split(".")[-1]
100class ShrinkyImage:
101 """does all the things"""
103 def __init__(self, source_path: Path) -> None:
104 """loads the image"""
105 self.source_path = source_path
106 try:
107 self.image = Image.open(source_path.open("rb"))
108 except UnidentifiedImageError as image_error:
109 logger.error("Pillow can't handle the file '{}', bailing.", source_path)
110 logger.error(image_error)
111 raise image_error
113 logger.debug("Dims: {}x{}", self.image.width, self.image.height)
114 logger.info(
115 "Original file size: {}", os.stat(self.source_path.resolve()).st_size
116 )
118 def resize_image(
119 self,
120 new_width: int,
121 new_height: int,
122 source_image: Optional[Image.Image] = None,
123 ) -> Image.Image:
124 """resizes an image, doesn't modify the source image"""
125 if source_image is None:
126 source_image = self.image
128 tmpimage = source_image.copy()
130 if source_image.width > new_width or source_image.height > new_height:
131 logger.debug(
132 "Thumbnailing from {}x{} to {}x{}",
133 source_image.width,
134 source_image.height,
135 new_width,
136 new_height,
137 )
138 tmpimage.thumbnail((new_width, new_height))
139 return tmpimage
141 def write_image(self, output_filename: Path, quality: int = -1) -> bool:
142 """writes the file to disk"""
144 try:
145 file_extension = get_extension(output_filename).lower()
146 except ValueError as extension_error:
147 logger.error(extension_error)
148 raise extension_error
150 args: Dict[str, Union[str, int]] = {}
152 if file_extension in ("jpg", "jpeg"):
153 # set jpeg quality
154 if quality is not None and quality >= 0:
155 args["quality"] = quality
157 if self.image.mode != "RGB":
158 logger.debug("Image is not RGB: {}", self.image.mode)
159 self.image = self.image.convert("RGB")
161 with output_filename.open("wb") as output_image:
162 logger.info("Writing {}", output_image.name)
163 self.image.save(output_image, **args) # type: ignore
165 new_size = os.stat(output_filename.resolve()).st_size
166 logger.info("New size: {}", new_size)
167 return True
170def setup_logging(
171 logger_object: Any,
172 debug: bool,
173) -> None:
174 """sets up loguru"""
175 logger_object.remove()
176 format_string = "<level>{level: <8}</level> - <level>{message}</level>"
177 if debug:
178 level = "DEBUG"
180 else:
181 level = "INFO"
182 logger_object.add(sink=sys.stdout, format=format_string, level=level)
185@click.command()
186@click.argument("filename", type=click.Path(exists=False, path_type=Path))
187@click.option(
188 "-o",
189 "--output",
190 type=click.Path(exists=False, dir_okay=False, resolve_path=True, path_type=Path),
191)
192@click.option("-t", "--output-type", help="New file type (eg jpg, png avif etc)")
193@click.option("-g", "--geometry", help="Geometry, 1x1, 1x, x1 etc.")
194@click.option("-q", "--quality", type=int, help="If JPEG, set quality.")
195@click.option("-f", "--force", is_flag=True, help="Overwrite destination")
196@click.option(
197 "--delete-source",
198 is_flag=True,
199 default=False,
200 help="Delete the source file once done",
201)
202@click.option("--debug", "-d", is_flag=True, help="Enable debug logging")
203def cli(
204 filename: Path = Path("~/"),
205 output: Optional[Path] = None,
206 output_type: Optional[str] = None,
207 force: bool = False,
208 quality: int = -1,
209 geometry: Optional[str] = None,
210 delete_source: bool = False,
211 debug: bool = False,
212) -> bool:
213 """Shrinky shrinks images in a way I like"""
215 setup_logging(logger, debug)
217 image_dimensions = set_geometry(geometry_value=geometry)
219 if output_type is not None:
220 if output_type.lower() not in VALID_OUTPUT_TYPES:
221 logger.error(
222 "Invalid output type {}, valid types are: {}",
223 output_type,
224 VALID_OUTPUT_TYPES,
225 )
226 sys.exit(1)
227 else:
228 logger.debug("Output type is {}", output_type.lower())
230 if output is None or output_type is not None:
231 output = new_filename(filename, output_type)
233 if output.exists() and not force:
234 logger.error("{} already exists, bailing", output.resolve())
235 sys.exit(1)
237 if get_extension(output).lower() not in VALID_OUTPUT_TYPES:
238 logger.error(
239 "Invalid output type {}, valid types are: {}",
240 get_extension(output),
241 VALID_OUTPUT_TYPES,
242 )
243 sys.exit(1)
245 original_file = Path(filename).resolve()
246 if not original_file.exists():
247 logger.error("Can't find {}, bailing", original_file)
248 sys.exit(1)
250 # hacky workaround for the avif extension - https://github.com/python-pillow/Pillow/pull/5201
251 if (
252 get_extension(output).lower() == "avif"
253 and "avif" not in Image.registered_extensions()
254 ):
255 import pillow_avif # type: ignore # noqa: F401,unused-import,import-outside-toplevel
257 image = ShrinkyImage(original_file)
259 # resize and store
260 image.image = image.resize_image(*image_dimensions)
262 image.write_image(
263 output,
264 quality=quality,
265 )
267 if delete_source:
268 logger.info("Please confirm you want to remove {} (y/N):", original_file)
269 if input().strip().lower() == "y":
270 original_file.unlink()
271 logger.info("Deleted {}", original_file)
272 else:
273 logger.info("Cancelled at user's request.")
274 return True