Refactor sphinx.util.images.guess_mimetype

Author: AA-Turner

sphinx/util/images.py

@@ -3,27 +3,29 @@
 from __future__ import annotations
 
 import base64
-import imghdr
-from collections import OrderedDict
 from os import path
-from typing import IO, BinaryIO, NamedTuple
+from typing import TYPE_CHECKING, NamedTuple
 
 import imagesize
 
+if TYPE_CHECKING:
+    from os import PathLike
+
 try:
     from PIL import Image
 except ImportError:
     Image = None
 
-mime_suffixes = OrderedDict([
-    ('.gif', 'image/gif'),
-    ('.jpg', 'image/jpeg'),
-    ('.png', 'image/png'),
-    ('.pdf', 'application/pdf'),
-    ('.svg', 'image/svg+xml'),
-    ('.svgz', 'image/svg+xml'),
-    ('.ai', 'application/illustrator'),
-])
+mime_suffixes = {
+    '.gif': 'image/gif',
+    '.jpg': 'image/jpeg',
+    '.png': 'image/png',
+    '.pdf': 'application/pdf',
+    '.svg': 'image/svg+xml',
+    '.svgz': 'image/svg+xml',
+    '.ai': 'application/illustrator'
+}
+_suffix_from_mime = {v: k for k, v in reversed(mime_suffixes.items())}
 
 
 class DataURI(NamedTuple):
@@ -49,31 +51,25 @@ def get_image_size(filename: str) -> tuple[int, int] | None:
         return None
 
 
-def guess_mimetype_for_stream(stream: IO, default: str | None = None) -> str | None:
-    imgtype = imghdr.what(stream)
-    if imgtype:
-        return 'image/' + imgtype
-    else:
-        return default
-
-
-def guess_mimetype(filename: str = '', default: str | None = None) -> str | None:
-    _, ext = path.splitext(filename.lower())
+def guess_mimetype(
+    filename: PathLike[str] | str = '',
+    default: str | None = None,
+) -> str | None:
+    ext = path.splitext(filename)[1].lower()
     if ext in mime_suffixes:
         return mime_suffixes[ext]
-    elif path.exists(filename):
-        with open(filename, 'rb') as f:
-            return guess_mimetype_for_stream(f, default=default)
-
+    if path.exists(filename):
+        try:
+            imgtype = _image_type_from_file(filename)
+        except ValueError:
+            pass
+        else:
+            return 'image/' + imgtype
     return default
 
 
 def get_image_extension(mimetype: str) -> str | None:
-    for ext, _mimetype in mime_suffixes.items():
-        if mimetype == _mimetype:
-            return ext
-
-    return None
+    return _suffix_from_mime.get(mimetype)
 
 
 def parse_data_uri(uri: str) -> DataURI | None:
@@ -97,17 +93,43 @@ def parse_data_uri(uri: str) -> DataURI | None:
     return DataURI(mimetype, charset, image_data)
 
 
-def test_svg(h: bytes, f: BinaryIO | None) -> str | None:
-    """An additional imghdr library helper; test the header is SVG's or not."""
-    try:
-        if '<svg' in h.decode().lower():
-            return 'svg+xml'
-    except UnicodeDecodeError:
-        pass
+def _image_type_from_file(filename: PathLike[str] | str) -> str:
+    with open(filename, 'rb') as f:
+        header = f.read(32)  # 32 bytes
+
+    # Bitmap
+    # https://en.wikipedia.org/wiki/BMP_file_format#Bitmap_file_header
+    if header.startswith(b'BM'):
+        return 'bmp'
+
+    # GIF
+    # https://en.wikipedia.org/wiki/GIF#File_format
+    if header.startswith((b'GIF87a', b'GIF89a')):
+        return 'gif'
+
+    # JPEG data
+    # https://en.wikipedia.org/wiki/JPEG_File_Interchange_Format#File_format_structure
+    if header.startswith(b'\xFF\xD8'):
+        return 'jpeg'
+
+    # Portable Network Graphics
+    # https://en.wikipedia.org/wiki/PNG#File_header
+    if header.startswith(b'\x89PNG\r\n\x1A\n'):
+        return 'png'
+
+    # Scalable Vector Graphics
+    # https://svgwg.org/svg2-draft/struct.html
+    if b'<svg' in header.lower():
+        return 'svg+xml'
 
-    return None
+    # TIFF
+    # https://en.wikipedia.org/wiki/TIFF#Byte_order
+    if header.startswith((b'MM', b'II')):
+        return 'tiff'
 
+    # WebP
+    # https://en.wikipedia.org/wiki/WebP#Technology
+    if header.startswith(b'RIFF') and header[8:12] == b'WEBP':
+        return 'webp'
 
-# install test_svg() to imghdr
-# refs: https://docs.python.org/3.9/library/imghdr.html#imghdr.tests
-imghdr.tests.append(test_svg)
+    raise ValueError('Could not detect image type!')