fix #133 webp generate-thumbnail failed (#134)

lucemia · web-flow · commit 2dfa2a72d677 · 2024-06-27T22:33:01.000+08:00
* fix #133 webp generate-thumbnail failed * add uniitest * fix
diff --git a/.gitignore b/.gitignore
@@ -158,3 +158,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+.DS_Store
diff --git a/src/ffmpeg_media_type/exceptions.py b/src/ffmpeg_media_type/exceptions.py
@@ -1,4 +1,10 @@
-class FfmpegMediaTypeError(Exception):
+class FFmpegMediaTypeError(Exception):
     """
     Base class for exceptions in this module.
     """
+
+
+class FFMpegMediaCorruptedError(FFmpegMediaTypeError):
+    """
+    Exception raised when the media file is corrupted.
+    """
diff --git a/src/ffmpeg_media_type/info.py b/src/ffmpeg_media_type/info.py
@@ -2,6 +2,8 @@
 from pathlib import Path
 from urllib.parse import urlparse
 
+from ffmpeg_media_type.exceptions import FFMpegMediaCorruptedError
+
 from .schema import FFMpegSupport, MediaInfo
 from .utils.cache import load
 from .utils.ffprobe import ffprobe
@@ -52,7 +54,8 @@ def detect(uri: str | Path) -> MediaInfo:
         the media type information
 
     Raises:
-        FfmpegMediaTypeError: If the ffmpeg command fails.
+        FFmpegMediaTypeError: If the ffmpeg command fails.
+        FFMpegMediaCorruptedError: If the media file is corrupted.
     """
     uri = str(uri)
     info = ffprobe(uri)
@@ -63,7 +66,8 @@ def detect(uri: str | Path) -> MediaInfo:
 
     # NOTE: handle ffmpeg's image compatibility
     if format_name == "image2":
-        assert info.streams[0].codec_name
+        if not info.streams[0].codec_name:
+            raise FFMpegMediaCorruptedError(f"Corrupted image file {uri}")
         format_name = info.streams[0].codec_name
 
     # NOTE: detect file extension
@@ -81,8 +85,11 @@ def detect(uri: str | Path) -> MediaInfo:
     else:
         suggest_ext = None
 
-    # NOTE: we classify gif as image
+    # NOTE: we classify gif and mjpeg as imageg
     if not duration or format_name in ("gif", "mjpeg"):
+        if not (info.streams[0].width and info.streams[0].height):
+            raise FFMpegMediaCorruptedError(f"Corrupted image file {uri}")
+
         return MediaInfo(
             type="image",
             width=info.streams[0].width or 0,
@@ -99,6 +106,9 @@ def detect(uri: str | Path) -> MediaInfo:
             width = stream.width
             height = stream.height
 
+            if not (width and height):
+                raise FFMpegMediaCorruptedError(f"Corrupted image file {uri}")
+
             return MediaInfo(
                 type="video",
                 width=width or 0,
diff --git a/src/ffmpeg_media_type/test_data/133-raise-exception/tmpz1wcuugh.webp b/src/ffmpeg_media_type/test_data/133-raise-exception/tmpz1wcuugh.webp
@@ -0,0 +1,22 @@
+<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
+<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="ja" lang="ja">
+<head>
+<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
+<meta http-equiv="Content-Style-Type" content="text/css" />
+<meta http-equiv="Content-Script-Type" content="text/javascript" />
+
+<meta http-equiv="refresh" content="0;URL=http://web.gekisaka.jp/frontpage?from=cf403">
+<title>ERROR: The request could not be satisfied</title>
+</head>
+<body>
+
+<h1>ERROR</h1>
+<h2>The request could not be satisfied.</h2>
+<hr noshade="" size="1px">Request blocked.
+<br clear="all">
+<hr noshade="" size="1px">
+<pre>Generated by cloudfront (CloudFront)
+</pre>
+
+</body>
+</html>
diff --git a/src/ffmpeg_media_type/tests/__snapshots__/test_info.ambr b/src/ffmpeg_media_type/tests/__snapshots__/test_info.ambr
@@ -0,0 +1,4 @@
+# serializer version: 1
+# name: test_detect[tmpz1wcuugh.webp]
+  "<ExceptionInfo FFMpegMediaCorruptedError('Corrupted image file src/ffmpeg_media_type/test_data/133-raise-exception/tmpz1wcuugh.webp') tblen=2>"
+# ---
diff --git a/src/ffmpeg_media_type/tests/test_info.py b/src/ffmpeg_media_type/tests/test_info.py
@@ -7,15 +7,22 @@
 from syrupy.extensions.json import JSONSnapshotExtension
 
 from ..conftest import sample_test_media_files
+from ..exceptions import FFmpegMediaTypeError
 from ..info import detect, generate_thumbnail
 
 
 @pytest.mark.parametrize("case", sample_test_media_files())
 def test_detect(case: Path, snapshot: SnapshotAssertion) -> None:
     # NOTE:
     # because we use docker to run ffmpeg, the input path need to be relative to the current working directory
-    info = detect(str(case.relative_to(Path.cwd())))
-    snapshot(extension_class=JSONSnapshotExtension) == asdict(info)
 
-    if info.type in ("video", "image"):
-        assert os.path.exists(generate_thumbnail(str(case.relative_to(Path.cwd())), "tmp.png"))
+    if "raise-exception" in str(case):
+        with pytest.raises(FFmpegMediaTypeError) as e:
+            info = detect(str(case.relative_to(Path.cwd())))
+        assert snapshot == str(e)
+    else:
+        info = detect(str(case.relative_to(Path.cwd())))
+        assert snapshot(extension_class=JSONSnapshotExtension) == asdict(info)
+
+        if info.type in ("video", "image"):
+            assert os.path.exists(generate_thumbnail(str(case.relative_to(Path.cwd())), "tmp.png"))
diff --git a/src/ffmpeg_media_type/utils/ffprobe.py b/src/ffmpeg_media_type/utils/ffprobe.py
@@ -18,7 +18,7 @@ def ffprobe(input_url: str | Path) -> FFProbeInfo:
         the media information
 
     Raises:
-        FfmpegMediaTypeError: If the FFprobe command fails.
+        FFmpegMediaTypeError: If the FFprobe command fails.
     """
 
     input_url = hotfix_animate_webp(input_url)
diff --git a/src/ffmpeg_media_type/utils/hotfix_webp.py b/src/ffmpeg_media_type/utils/hotfix_webp.py
@@ -78,7 +78,7 @@ def is_webp_animated(file_path: str) -> bool:
         True if the WebP file is animated, False otherwise.
 
     Raises:
-        FfmpegMediaTypeError: If the webpmux command fails.
+        FFmpegMediaTypeError: If the webpmux command fails.
     """
 
     # Running the webpmux command to get information about the WebP file
@@ -103,7 +103,7 @@ def extract_animated_webp_frame(uri: str) -> str:
         The URI of the fixed webp file.
 
     Raises:
-        FfmpegMediaTypeError: If the webpmux command fails.
+        FFmpegMediaTypeError: If the webpmux command fails.
 
     Notes:
         Will raise Exception if the webp file is not animated.
@@ -125,7 +125,7 @@ def is_webp_need_fix(uri: str | Path) -> bool:
         True if the WebP file is animated, False otherwise.
 
     Raises:
-        FfmpegMediaTypeError: If the ffprobe command fails.
+        FFmpegMediaTypeError: If the ffprobe command fails.
     """
 
     ffprobe_cmd = ["ffprobe"] + [
@@ -162,7 +162,7 @@ def hotfix_animate_webp(uri: str | Path) -> str:
         The URI of the fixed WebP file if it was animated, otherwise the original URI.
 
     Raises:
-        FfmpegMediaTypeError: If the webpmux command fails.
+        FFmpegMediaTypeError: If the webpmux command fails.
 
     Notes:
         Some WebP files (e.g. animated WebP files) may not be supported by ffmpeg.
diff --git a/src/ffmpeg_media_type/utils/shell.py b/src/ffmpeg_media_type/utils/shell.py
@@ -2,7 +2,7 @@
 from shlex import join
 from tempfile import NamedTemporaryFile
 
-from ..exceptions import FfmpegMediaTypeError
+from ..exceptions import FFmpegMediaTypeError
 
 
 def call(cmds: list[str]) -> str:
@@ -13,15 +13,15 @@ def call(cmds: list[str]) -> str:
         cmds: List of command and arguments.
 
     Raises:
-        FfmpegMediaTypeError: If the command fails.
+        FFmpegMediaTypeError: If the command fails.
 
     Returns:
         stdout of the command.
     """
     try:
         r = subprocess.run(cmds, stdout=subprocess.PIPE, check=True)
     except subprocess.CalledProcessError as e:
-        raise FfmpegMediaTypeError(f"command failed: {join(e.cmd)}")
+        raise FFmpegMediaTypeError(f"command failed: {join(e.cmd)}")
 
     return r.stdout.decode("utf-8")
 
diff --git a/src/ffmpeg_media_type/utils/tests/__snapshots__/test_ffprobe/test_ffprobe_file[tmpz1wcuugh.webp].json b/src/ffmpeg_media_type/utils/tests/__snapshots__/test_ffprobe/test_ffprobe_file[tmpz1wcuugh.webp].json
@@ -0,0 +1,26 @@
+{
+  "format": {
+    "duration": "0.040000",
+    "format_long_name": "image2 sequence",
+    "format_name": "image2",
+    "probe_score": 50,
+    "size": "790",
+    "start_time": "0.000000"
+  },
+  "streams": [
+    {
+      "codec_long_name": "WebP",
+      "codec_name": "webp",
+      "codec_type": "video",
+      "height": 0,
+      "index": 0,
+      "pix_fmt": null,
+      "profile": null,
+      "r_frame_rate": "25/1",
+      "tags": {
+        "rotate": 0
+      },
+      "width": 0
+    }
+  ]
+}
diff --git a/src/ffmpeg_media_type/utils/tests/__snapshots__/test_thumbnail.ambr b/src/ffmpeg_media_type/utils/tests/__snapshots__/test_thumbnail.ambr
@@ -50,3 +50,6 @@
 # name: test_thumbnail[silent_video.mp4]
   True
 # ---
+# name: test_thumbnail[tmpz1wcuugh.webp]
+  False
+# ---
diff --git a/src/ffmpeg_media_type/utils/tests/test_hotfix_webp.py b/src/ffmpeg_media_type/utils/tests/test_hotfix_webp.py
@@ -4,7 +4,7 @@
 from syrupy.assertion import SnapshotAssertion
 
 from ...conftest import sample_test_media_files
-from ...exceptions import FfmpegMediaTypeError
+from ...exceptions import FFmpegMediaTypeError
 from ..hotfix_webp import extract_animated_webp_frame, hotfix_animate_webp, is_webp_animated
 
 
@@ -18,7 +18,7 @@ def test_extract_animated_webp_frame(case: Path, snapshot: SnapshotAssertion) ->
     try:
         extract_animated_webp_frame(str(case))
         assert snapshot == True
-    except FfmpegMediaTypeError:
+    except FFmpegMediaTypeError:
         assert snapshot == False
 
 
diff --git a/src/ffmpeg_media_type/utils/tests/test_thumbnail.py b/src/ffmpeg_media_type/utils/tests/test_thumbnail.py
@@ -4,7 +4,7 @@
 from syrupy.assertion import SnapshotAssertion
 
 from ...conftest import sample_test_media_files
-from ...exceptions import FfmpegMediaTypeError
+from ...exceptions import FFmpegMediaTypeError
 from ..thumbnail import generate_thumbnail
 
 
@@ -13,5 +13,5 @@ def test_thumbnail(case: Path, snapshot: SnapshotAssertion) -> None:
     try:
         generate_thumbnail(case)
         assert snapshot == True
-    except FfmpegMediaTypeError:
+    except FFmpegMediaTypeError:
         assert snapshot == False
diff --git a/src/ffmpeg_media_type/utils/thumbnail.py b/src/ffmpeg_media_type/utils/thumbnail.py
@@ -5,7 +5,12 @@
 
 
 def generate_thumbnail(
-    video_path: str | Path, suffix: str = ".png", *, width: int = 320, height: int = -1, time_offset: float = 0
+    video_path: str | Path,
+    suffix: str = ".png",
+    *,
+    width: int = 320,
+    height: int = -1,
+    time_offset: float = 0,
 ) -> str:
     """
     Generate a thumbnail from a video file at a specified time offset.
@@ -18,7 +23,7 @@ def generate_thumbnail(
         time_offset: the time offset in seconds to generate the thumbnail
 
     Raises:
-        FfmpegMediaTypeError: If the ffmpeg command fails.
+        FFmpegMediaTypeError: If the ffmpeg command fails.
 
     Returns:
         the path to the generated thumbnail
