Add support for probing SVG images

openzim · Jul 11, 2024 · 0037caa · 0037caa
1 parent 02be5e4
commit 0037caa
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 11 deletions.
diff --git a/src/zimscraperlib/image/probing.py b/src/zimscraperlib/image/probing.py
@@ -12,6 +12,8 @@
 import colorthief
 import PIL.Image
 
+from zimscraperlib.filesystem import get_content_mimetype, get_file_mimetype
+
 
 def get_colors(
     src: pathlib.Path, *, use_palette: bool | None = True
@@ -59,8 +61,23 @@ def format_for(
 ) -> str | None:
     """Pillow format of a given filename, either Pillow-detected or from suffix"""
     if not from_suffix:
-        with PIL.Image.open(src) as img:
-            return img.format
+        try:
+            with PIL.Image.open(src) as img:
+                return img.format
+        except PIL.UnidentifiedImageError:
+            # Fallback based on mimetype for SVG which are not supported by PIL
+            if (
+                isinstance(src, pathlib.Path)
+                and get_file_mimetype(src) == "image/svg+xml"
+            ):
+                return "SVG"
+            elif (
+                isinstance(src, io.BytesIO)
+                and get_content_mimetype(src.getvalue()) == "image/svg+xml"
+            ):
+                return "SVG"
+            else:
+                raise
 
     if not isinstance(src, pathlib.Path):
         raise ValueError(
@@ -70,8 +87,11 @@ def format_for(
     from PIL.Image import EXTENSION as PIL_FMT_EXTENSION
     from PIL.Image import init as init_pil
 
-    init_pil()
-    return PIL_FMT_EXTENSION[src.suffix] if src.suffix in PIL_FMT_EXTENSION else None
+    init_pil()  # populate the PIL_FMT_EXTENSION dictionary
+
+    known_extensions = {".svg": "SVG"}
+    known_extensions.update(PIL_FMT_EXTENSION)
+    return known_extensions[src.suffix] if src.suffix in known_extensions else None
 
 
 def is_valid_image(

diff --git a/tests/image/test_image.py b/tests/image/test_image.py
@@ -68,8 +68,15 @@ def get_src_dst(
     jpg_image: pathlib.Path | None = None,
     gif_image: pathlib.Path | None = None,
     webp_image: pathlib.Path | None = None,
+    svg_image: pathlib.Path | None = None,
 ) -> tuple[pathlib.Path, pathlib.Path]:
-    options = {"png": png_image, "jpg": jpg_image, "webp": webp_image, "gif": gif_image}
+    options = {
+        "png": png_image,
+        "jpg": jpg_image,
+        "webp": webp_image,
+        "gif": gif_image,
+        "svg": svg_image,
+    }
     if fmt not in options:
         raise LookupError(f"Unsupported fmt passed: {fmt}")
     src = options[fmt]
@@ -616,10 +623,10 @@ def test_ensure_matches(webp_image):
 
 @pytest.mark.parametrize(
     "fmt,expected",
-    [("png", "PNG"), ("jpg", "JPEG"), ("gif", "GIF"), ("webp", "WEBP")],
+    [("png", "PNG"), ("jpg", "JPEG"), ("gif", "GIF"), ("webp", "WEBP"), ("svg", "SVG")],
 )
 def test_format_for_real_images_suffix(
-    png_image, jpg_image, gif_image, webp_image, tmp_path, fmt, expected
+    png_image, jpg_image, gif_image, webp_image, svg_image, tmp_path, fmt, expected
 ):
     src, _ = get_src_dst(
         tmp_path,
@@ -628,16 +635,17 @@ def test_format_for_real_images_suffix(
         jpg_image=jpg_image,
         gif_image=gif_image,
         webp_image=webp_image,
+        svg_image=svg_image,
     )
     assert format_for(src) == expected
 
 
 @pytest.mark.parametrize(
     "fmt,expected",
-    [("png", "PNG"), ("jpg", "JPEG"), ("gif", "GIF"), ("webp", "WEBP")],
+    [("png", "PNG"), ("jpg", "JPEG"), ("gif", "GIF"), ("webp", "WEBP"), ("svg", "SVG")],
 )
 def test_format_for_real_images_content_path(
-    png_image, jpg_image, gif_image, webp_image, tmp_path, fmt, expected
+    png_image, jpg_image, gif_image, webp_image, svg_image, tmp_path, fmt, expected
 ):
     src, _ = get_src_dst(
         tmp_path,
@@ -646,16 +654,17 @@ def test_format_for_real_images_content_path(
         jpg_image=jpg_image,
         gif_image=gif_image,
         webp_image=webp_image,
+        svg_image=svg_image,
     )
     assert format_for(src, from_suffix=False) == expected
 
 
 @pytest.mark.parametrize(
     "fmt,expected",
-    [("png", "PNG"), ("jpg", "JPEG"), ("gif", "GIF"), ("webp", "WEBP")],
+    [("png", "PNG"), ("jpg", "JPEG"), ("gif", "GIF"), ("webp", "WEBP"), ("svg", "SVG")],
 )
 def test_format_for_real_images_content_bytes(
-    png_image, jpg_image, gif_image, webp_image, tmp_path, fmt, expected
+    png_image, jpg_image, gif_image, webp_image, svg_image, tmp_path, fmt, expected
 ):
     src, _ = get_src_dst(
         tmp_path,
@@ -664,6 +673,7 @@ def test_format_for_real_images_content_bytes(
         jpg_image=jpg_image,
         gif_image=gif_image,
         webp_image=webp_image,
+        svg_image=svg_image,
     )
     assert format_for(io.BytesIO(src.read_bytes()), from_suffix=False) == expected
 
@@ -675,6 +685,7 @@ def test_format_for_real_images_content_bytes(
         ("image.jpg", "JPEG"),
         ("image.gif", "GIF"),
         ("image.webp", "WEBP"),
+        ("image.svg", "SVG"),
         ("image.raster", None),
     ],
 )