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

* fix #133 webp generate-thumbnail failed

* add uniitest

* fix

Author: lucemia

.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

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.
+    """

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,

src/ffmpeg_media_type/test_data/133-raise-exception/tmpz1wcuugh.webp

@@ -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>"
+# ---

src/ffmpeg_media_type/tests/__snapshots__/test_info.ambr

@@ -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"))

src/ffmpeg_media_type/tests/test_info.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)

src/ffmpeg_media_type/utils/ffprobe.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.

src/ffmpeg_media_type/utils/hotfix_webp.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")
 

src/ffmpeg_media_type/utils/shell.py

@@ -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
+    }
+  ]
+}