Skip to content

Commit 78fa538

Browse files
committed
Make image optimization methods stricter with options types
1 parent 8a98b12 commit 78fa538

File tree

3 files changed

+497
-216
lines changed

3 files changed

+497
-216
lines changed

src/zimscraperlib/image/optimization.py

Lines changed: 147 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,11 @@
1818
can still run on default settings which give
1919
a bit less size than the original images but maintain a high quality. """
2020

21-
import functools
2221
import io
2322
import os
2423
import pathlib
2524
import subprocess
26-
from collections.abc import Callable
27-
from typing import Any
25+
from dataclasses import dataclass
2826

2927
import piexif # pyright: ignore[reportMissingTypeStubs]
3028
from optimize_images.img_aux_processing import ( # pyright: ignore[reportMissingTypeStubs]
@@ -54,18 +52,9 @@ def ensure_matches(
5452
raise ValueError(f"{src} is not of format {fmt}")
5553

5654

57-
def optimize_png(
58-
src: pathlib.Path | io.BytesIO,
59-
dst: pathlib.Path | io.BytesIO | None = None,
60-
max_colors: int = 256,
61-
background_color: tuple[int, int, int] = (255, 255, 255),
62-
*,
63-
reduce_colors: bool | None = False,
64-
fast_mode: bool | None = True,
65-
remove_transparency: bool | None = False,
66-
**_: Any,
67-
) -> pathlib.Path | io.BytesIO:
68-
"""method to optimize PNG files using a pure python external optimizer
55+
@dataclass
56+
class OptimizePngOptions:
57+
"""Dataclass holding PNG optimization options
6958
7059
Arguments:
7160
reduce_colors: Whether to reduce colors using adaptive color pallette (boolean)
@@ -79,20 +68,38 @@ def optimize_png(
7968
values: True | False
8069
background_color: Background color if remove_transparency is True (tuple
8170
containing RGB values)
82-
values: (255, 255, 255) | (221, 121, 108) | (XX, YY, ZZ)"""
71+
values: (255, 255, 255) | (221, 121, 108) | (XX, YY, ZZ)
72+
"""
73+
74+
max_colors: int = 256
75+
background_color: tuple[int, int, int] = (255, 255, 255)
76+
reduce_colors: bool | None = False
77+
fast_mode: bool | None = True
78+
remove_transparency: bool | None = False
79+
80+
81+
def optimize_png(
82+
src: pathlib.Path | io.BytesIO,
83+
dst: pathlib.Path | io.BytesIO | None = None,
84+
options: OptimizePngOptions | None = None,
85+
) -> pathlib.Path | io.BytesIO:
86+
"""method to optimize PNG files using a pure python external optimizer"""
8387

8488
ensure_matches(src, "PNG")
8589

8690
img = Image.open(src)
8791

88-
if remove_transparency:
89-
img = remove_alpha(img, background_color)
92+
if options is None:
93+
options = OptimizePngOptions()
94+
95+
if options.remove_transparency:
96+
img = remove_alpha(img, options.background_color)
9097

91-
if reduce_colors:
92-
img, __, __ = do_reduce_colors(img, max_colors)
98+
if options.reduce_colors:
99+
img, _, _ = do_reduce_colors(img, options.max_colors)
93100

94-
if not fast_mode and img.mode == "P":
95-
img, __ = rebuild_palette(img)
101+
if not options.fast_mode and img.mode == "P":
102+
img, _ = rebuild_palette(img)
96103

97104
if dst is None:
98105
dst = io.BytesIO()
@@ -102,16 +109,9 @@ def optimize_png(
102109
return dst
103110

104111

105-
def optimize_jpeg(
106-
src: pathlib.Path | io.BytesIO,
107-
dst: pathlib.Path | io.BytesIO | None = None,
108-
quality: int | None = 85,
109-
*,
110-
fast_mode: bool | None = True,
111-
keep_exif: bool | None = True,
112-
**_: Any,
113-
) -> pathlib.Path | io.BytesIO:
114-
"""method to optimize JPEG files using a pure python external optimizer
112+
@dataclass
113+
class OptimizeJpgOptions:
114+
"""Dataclass holding JPG optimization options
115115
116116
Arguments:
117117
quality: JPEG quality (integer between 1 and 100)
@@ -120,7 +120,23 @@ def optimize_jpeg(
120120
values: True | False
121121
fast_mode: Use the supplied quality value. If turned off, optimizer will
122122
get dynamic quality value to ensure better compression
123-
values: True | False"""
123+
values: True | False
124+
"""
125+
126+
quality: int | None = 85
127+
fast_mode: bool | None = True
128+
keep_exif: bool | None = True
129+
130+
131+
def optimize_jpeg(
132+
src: pathlib.Path | io.BytesIO,
133+
dst: pathlib.Path | io.BytesIO | None = None,
134+
options: OptimizeJpgOptions | None = None,
135+
) -> pathlib.Path | io.BytesIO:
136+
"""method to optimize JPEG files using a pure python external optimizer"""
137+
138+
if options is None:
139+
options = OptimizeJpgOptions()
124140

125141
ensure_matches(src, "JPEG")
126142

@@ -146,10 +162,10 @@ def optimize_jpeg(
146162
# only use progressive if file size is bigger
147163
use_progressive_jpg = orig_size > 10240 # 10KiB # noqa: PLR2004
148164

149-
if fast_mode:
150-
quality_setting = quality
165+
if options.fast_mode:
166+
quality_setting = options.quality
151167
else:
152-
quality_setting, __ = jpeg_dynamic_quality(img)
168+
quality_setting, _ = jpeg_dynamic_quality(img)
153169

154170
if dst is None:
155171
dst = io.BytesIO()
@@ -165,7 +181,7 @@ def optimize_jpeg(
165181
if isinstance(dst, io.BytesIO):
166182
dst.seek(0)
167183

168-
if keep_exif and had_exif:
184+
if options.keep_exif and had_exif:
169185
piexif.transplant( # pyright: ignore[reportUnknownMemberType]
170186
exif_src=(
171187
str(src.resolve()) if isinstance(src, pathlib.Path) else src.getvalue()
@@ -179,16 +195,9 @@ def optimize_jpeg(
179195
return dst
180196

181197

182-
def optimize_webp(
183-
src: pathlib.Path | io.BytesIO,
184-
dst: pathlib.Path | io.BytesIO | None = None,
185-
quality: int | None = 60,
186-
method: int | None = 6,
187-
*,
188-
lossless: bool | None = False,
189-
**_: Any,
190-
) -> pathlib.Path | io.BytesIO:
191-
"""method to optimize WebP using Pillow options
198+
@dataclass
199+
class OptimizeWebpOptions:
200+
"""Dataclass holding WebP optimization options
192201
193202
Arguments:
194203
lossless: Whether to use lossless compression (boolean);
@@ -201,13 +210,29 @@ def optimize_webp(
201210
values: 1 | 2 | 3 | 4 | 5 | 6
202211
203212
refer to the link for more details
204-
https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#webp"""
213+
https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#webp
214+
"""
215+
216+
quality: int | None = 60
217+
method: int | None = 6
218+
lossless: bool | None = False
219+
220+
221+
def optimize_webp(
222+
src: pathlib.Path | io.BytesIO,
223+
dst: pathlib.Path | io.BytesIO | None = None,
224+
options: OptimizeWebpOptions | None = None,
225+
) -> pathlib.Path | io.BytesIO:
226+
"""method to optimize WebP using Pillow options"""
227+
228+
if options is None:
229+
options = OptimizeWebpOptions()
205230

206231
ensure_matches(src, "WEBP")
207232
params: dict[str, bool | int | None] = {
208-
"lossless": lossless,
209-
"quality": quality,
210-
"method": method,
233+
"lossless": options.lossless,
234+
"quality": options.quality,
235+
"method": options.method,
211236
}
212237

213238
webp_image = Image.open(src)
@@ -230,18 +255,9 @@ def optimize_webp(
230255
return dst
231256

232257

233-
def optimize_gif(
234-
src: pathlib.Path,
235-
dst: pathlib.Path,
236-
optimize_level: int | None = 1,
237-
lossiness: int | None = None,
238-
max_colors: int | None = None,
239-
*,
240-
interlace: bool | None = True,
241-
no_extensions: bool | None = True,
242-
**_: Any,
243-
) -> pathlib.Path:
244-
"""method to optimize GIFs using gifsicle >= 1.92
258+
@dataclass
259+
class OptimizeGifOptions:
260+
"""Dataclass holding GIF optimization options
245261
246262
Arguments:
247263
optimize_level: Optimization level; higher values give better compression
@@ -258,21 +274,37 @@ def optimize_gif(
258274
(integer between 2 and 256)
259275
values: 2 | 86 | 128 | 256 | XX
260276
261-
refer to the link for more details - https://www.lcdf.org/gifsicle/man.html"""
277+
refer to the link for more details - https://www.lcdf.org/gifsicle/man.html
278+
"""
279+
280+
optimize_level: int | None = 1
281+
lossiness: int | None = None
282+
max_colors: int | None = None
283+
interlace: bool | None = True
284+
no_extensions: bool | None = True
285+
286+
287+
def optimize_gif(
288+
src: pathlib.Path, dst: pathlib.Path, options: OptimizeGifOptions | None = None
289+
) -> pathlib.Path:
290+
"""method to optimize GIFs using gifsicle >= 1.92"""
291+
292+
if options is None:
293+
options = OptimizeGifOptions()
262294

263295
ensure_matches(src, "GIF")
264296

265297
# use gifsicle
266298
args = ["/usr/bin/env", "gifsicle"]
267-
if optimize_level:
268-
args += [f"-O{optimize_level}"]
269-
if max_colors:
270-
args += ["--colors", str(max_colors)]
271-
if lossiness:
272-
args += [f"--lossy={lossiness}"]
273-
if no_extensions:
299+
if options.optimize_level:
300+
args += [f"-O{options.optimize_level}"]
301+
if options.max_colors:
302+
args += ["--colors", str(options.max_colors)]
303+
if options.lossiness:
304+
args += [f"--lossy={options.lossiness}"]
305+
if options.no_extensions:
274306
args += ["--no-extensions"]
275-
if interlace:
307+
if options.interlace:
276308
args += ["--interlace"]
277309
args += [str(src)]
278310
with open(dst, "w") as out_file:
@@ -287,13 +319,39 @@ def optimize_gif(
287319
return dst
288320

289321

322+
@dataclass
323+
class OptimizeOptions:
324+
"""Dataclass holding GIF optimization options for all supported formats"""
325+
326+
gif: OptimizeGifOptions
327+
webp: OptimizeWebpOptions
328+
jpg: OptimizeJpgOptions
329+
png: OptimizePngOptions
330+
331+
@classmethod
332+
def of(
333+
cls,
334+
gif: OptimizeGifOptions | None = None,
335+
webp: OptimizeWebpOptions | None = None,
336+
jpg: OptimizeJpgOptions | None = None,
337+
png: OptimizePngOptions | None = None,
338+
):
339+
"""Helper to override only few options from default value"""
340+
return OptimizeOptions(
341+
gif=gif or OptimizeGifOptions(),
342+
png=png or OptimizePngOptions(),
343+
webp=webp or OptimizeWebpOptions(),
344+
jpg=jpg or OptimizeJpgOptions(),
345+
)
346+
347+
290348
def optimize_image(
291349
src: pathlib.Path,
292350
dst: pathlib.Path,
351+
options: OptimizeOptions | None = None,
293352
*,
294353
delete_src: bool | None = False,
295354
convert: bool | str | None = False,
296-
**options: Any,
297355
):
298356
"""Optimize image, automatically selecting correct optimizer
299357
@@ -305,6 +363,9 @@ def optimize_image(
305363
True: convert to format implied by dst suffix
306364
"FMT": convert to format FMT (use Pillow names)"""
307365

366+
if options is None:
367+
options = OptimizeOptions.of()
368+
308369
src_format, dst_format = format_for(src, from_suffix=False), format_for(dst)
309370

310371
if src_format is None: # pragma: no cover
@@ -321,26 +382,20 @@ def optimize_image(
321382
else:
322383
src_img = pathlib.Path(src)
323384

324-
get_optimization_method(src_format)(src_img, dst, **options)
385+
src_format = src_format.lower()
386+
if src_format in ("jpg", "jpeg"):
387+
optimize_jpeg(src=src_img, dst=dst, options=options.jpg)
388+
elif src_format == "gif":
389+
optimize_gif(src=src_img, dst=dst, options=options.gif)
390+
elif src_format == "png":
391+
optimize_png(src=src_img, dst=dst, options=options.png)
392+
elif src_format == "webp":
393+
optimize_webp(src=src_img, dst=dst, options=options.webp)
394+
else:
395+
raise NotImplementedError(
396+
f"Image format '{src_format}' cannot yet be optimized"
397+
)
325398

326399
# delete src image if requested
327400
if delete_src and src.exists() and src.resolve() != dst.resolve():
328401
src.unlink()
329-
330-
331-
def get_optimization_method(fmt: str) -> Callable[..., Any]:
332-
"""Return the proper optimization method to call for a given image format"""
333-
334-
def raise_error(*_, orig_format: str):
335-
raise NotImplementedError(
336-
f"Image format '{orig_format}' cannot yet be optimized"
337-
)
338-
339-
fmt = fmt.lower().strip()
340-
return {
341-
"gif": optimize_gif,
342-
"jpg": optimize_jpeg,
343-
"jpeg": optimize_jpeg,
344-
"webp": optimize_webp,
345-
"png": optimize_png,
346-
}.get(fmt, functools.partial(raise_error, orig_format=fmt))

0 commit comments

Comments
 (0)