sticker-convert 2.3.0__py3-none-any.whl → 2.4.0__py3-none-any.whl

sticker_convert/utils/files/sanitize_filename.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+import unicodedata
+import re
+
+def sanitize_filename(filename: str) -> str:
+    # Based on https://gitlab.com/jplusplus/sanitize-filename/-/blob/master/sanitize_filename/sanitize_filename.py
+    # Replace illegal character with '_'
+    """Return a fairly safe version of the filename.
+
+    We don't limit ourselves to ascii, because we want to keep municipality
+    names, etc, but we do want to get rid of anything potentially harmful,
+    and make sure we do not exceed Windows filename length limits.
+    Hence a less safe blacklist, rather than a whitelist.
+    """
+    blacklist = ["\\", "/", ":", "*", "?", '"', "<", ">", "|", "\0"]
+    reserved = [
+        "CON", "PRN", "AUX", "NUL", "COM1", "COM2", "COM3", "COM4", "COM5",
+        "COM6", "COM7", "COM8", "COM9", "LPT1", "LPT2", "LPT3", "LPT4", "LPT5",
+        "LPT6", "LPT7", "LPT8", "LPT9",
+    ]  # Reserved words on Windows
+    filename = "".join(c if c not in blacklist else "_" for c in filename)
+    # Remove all charcters below code point 32
+    filename = "".join(c if 31 < ord(c) else "_" for c in filename)
+    filename = unicodedata.normalize("NFKD", filename)
+    filename = filename.rstrip(". ")  # Windows does not allow these at end
+    filename = filename.strip()
+    if all([x == "." for x in filename]):
+        filename = "__" + filename
+    if filename in reserved:
+        filename = "__" + filename
+    if len(filename) == 0:
+        filename = "__"
+    if len(filename) > 255:
+        parts = re.split(r"/|\\", filename)[-1].split(".")
+        if len(parts) > 1:
+            ext = "." + parts.pop()
+            filename = filename[: -len(ext)]
+        else:
+            ext = ""
+        if filename == "":
+            filename = "__"
+        if len(ext) > 254:
+            ext = ext[254:]
+        maxl = 255 - len(ext)
+        filename = filename[:maxl]
+        filename = filename + ext
+        # Re-check last character (if there was no extension)
+        filename = filename.rstrip(". ")
+        if len(filename) == 0:
+            filename = "__"
+    return filename

sticker_convert/utils/media/codec_info.py

@@ -1,161 +1,227 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 import os
-import mimetypes
+import mmap
 from typing import Optional
 
-import imageio.v3 as iio
-import av # type: ignore
-from av.codec.context import CodecContext # type: ignore
-from rlottie_python import LottieAnimation  # type: ignore
 from PIL import Image, UnidentifiedImageError
-import mmap
-
 
 class CodecInfo:
-    def __init__(self):
-        mimetypes.init()
-        vid_ext = []
-        for ext in mimetypes.types_map:
-            if mimetypes.types_map[ext].split("/")[0] == "video":
-                vid_ext.append(ext)
-        vid_ext.append(".webm")
-        vid_ext.append(".webp")
-        self.vid_ext = tuple(vid_ext)
+    def __init__(self, file: str):
+        self.file_ext = CodecInfo.get_file_ext(file)
+        self.fps, self.frames, self.duration = CodecInfo.get_file_fps_frames_duration(file)
+        self.codec = CodecInfo.get_file_codec(file)
+        self.res = CodecInfo.get_file_res(file)
+        self.is_animated = True if self.fps > 1 else False
 
     @staticmethod
-    def get_file_fps(file: str) -> float:
+    def get_file_fps_frames_duration(file: str) -> tuple[float, int, int]:
         file_ext = CodecInfo.get_file_ext(file)
 
-        if file_ext in ".tgs":
-            with LottieAnimation.from_tgs(file) as anim:
-                fps = anim.lottie_animation_get_framerate()
+        if file_ext == ".tgs":
+            fps, frames = CodecInfo._get_file_fps_frames_tgs(file)
+            duration = int(frames / fps * 1000)
         else:
             if file_ext == ".webp":
-                total_duration = 0
-                frames = 0
-                
-                with open(file, "r+b") as f:
-                    mm = mmap.mmap(f.fileno(), 0)
-                    while True:
-                        anmf_pos = mm.find(b"ANMF")
-                        if anmf_pos == -1:
-                            break
-                        mm.seek(anmf_pos + 20)
-                        frame_duration_32 = mm.read(4)
-                        frame_duration = frame_duration_32[:-1] + bytes(
-                            int(frame_duration_32[-1]) & 0b11111100
-                        )
-                        total_duration += int.from_bytes(frame_duration, "little")
-                        frames += 1
-                
-                if frames == 0:
-                    fps = 1
-                else:
-                    fps = frames / total_duration * 1000
+                frames, duration = CodecInfo._get_file_frames_duration_webp(file)
             elif file_ext in (".gif", ".apng", ".png"):
-                total_duration = 0
-                frames = len([* iio.imiter(file, plugin="pillow")])
+                frames, duration = CodecInfo._get_file_frames_duration_pillow(file)
+            else:
+                frames, duration = CodecInfo._get_file_frames_duration_av(file)
+
+            fps = frames / duration * 1000
+        
+        return fps, frames, duration
+
+    @staticmethod
+    def get_file_fps(file: str) -> float:
+        file_ext = CodecInfo.get_file_ext(file)
+
+        if file_ext == ".tgs":
+            return CodecInfo._get_file_fps_tgs(file)
+        elif file_ext == ".webp":
+            frames, duration = CodecInfo._get_file_frames_duration_webp(file)
+        elif file_ext in (".gif", ".apng", ".png"):
+            frames, duration = CodecInfo._get_file_frames_duration_pillow(file)
+        else:
+            frames, duration = CodecInfo._get_file_frames_duration_av(file, frames_to_iterate=10)
+            
+        return frames / duration * 1000
+    
+    @staticmethod
+    def get_file_frames(file: str, check_anim: bool = False) -> int:
+        # If check_anim is True, return value > 1 means the file is animated
+        file_ext = CodecInfo.get_file_ext(file)
+
+        if file_ext == ".tgs":
+            return CodecInfo._get_file_frames_tgs(file)
+        elif file_ext in (".gif", ".webp", ".png", ".apng"):
+            frames, _ = CodecInfo._get_file_frames_duration_pillow(file, frames_only=True)
+        else:
+            if check_anim == True:
+                frames_to_iterate = 2
+            else:
+                frames_to_iterate = None
+            frames, _ = CodecInfo._get_file_frames_duration_av(file, frames_only=True, frames_to_iterate=frames_to_iterate)
+
+        return frames
+
+    @staticmethod
+    def get_file_duration(file: str) -> int:
+        # Return duration in miliseconds
+        file_ext = CodecInfo.get_file_ext(file)
+
+        if file_ext == ".tgs":
+            fps, frames = CodecInfo._get_file_fps_frames_tgs(file)
+            duration = int(frames / fps * 1000)
+        elif file_ext == ".webp":
+            _, duration = CodecInfo._get_file_frames_duration_webp(file)
+        elif file_ext in (".gif", ".png", ".apng"):
+            _, duration = CodecInfo._get_file_frames_duration_pillow(file)
+        else:
+            _, duration = CodecInfo._get_file_frames_duration_av(file)
+        
+        return duration
+    
+    @staticmethod
+    def _get_file_fps_tgs(file: str) -> int:
+        from rlottie_python import LottieAnimation  # type: ignore
+
+        with LottieAnimation.from_tgs(file) as anim:
+            return anim.lottie_animation_get_framerate()
+
+    @staticmethod    
+    def _get_file_frames_tgs(file: str) -> int:
+        from rlottie_python import LottieAnimation  # type: ignore
+
+        with LottieAnimation.from_tgs(file) as anim:
+            return anim.lottie_animation_get_totalframe()
+    
+    @staticmethod
+    def _get_file_fps_frames_tgs(file: str) -> tuple[int, int]:
+        from rlottie_python import LottieAnimation  # type: ignore
+
+        with LottieAnimation.from_tgs(file) as anim:
+            fps = anim.lottie_animation_get_framerate()
+            frames = anim.lottie_animation_get_totalframe()
 
-                for frame in range(frames):
-                    metadata = iio.immeta(
-                        file, index=frame, plugin="pillow", exclude_applied=False
+        return fps, frames
+    
+    @staticmethod
+    def _get_file_frames_duration_pillow(file: str, frames_only: bool = False) -> tuple[int, int]:
+        total_duration = 0
+
+        with Image.open(file) as im:
+            if 'n_frames' in im.__dir__():
+                frames = im.n_frames
+                if frames_only == True:
+                    return frames, 1
+                for i in range(im.n_frames):
+                    im.seek(i)
+                    total_duration += im.info.get('duration', 1000)
+                return frames, total_duration
+            else:
+                return 1, 1
+        
+    @staticmethod
+    def _get_file_frames_duration_webp(file: str) -> tuple[int, int]:
+        total_duration = 0
+        frames = 0
+        
+        with open(file, "r+b") as f:
+            with mmap.mmap(f.fileno(), 0) as mm:
+                while True:
+                    anmf_pos = mm.find(b"ANMF")
+                    if anmf_pos == -1:
+                        break
+                    mm.seek(anmf_pos + 20)
+                    frame_duration_32 = mm.read(4)
+                    frame_duration = frame_duration_32[:-1] + bytes(
+                        int(frame_duration_32[-1]) & 0b11111100
                     )
-                    total_duration += metadata.get("duration", 1000)
-                
-                if frames == 0 or total_duration == 0:
-                    fps = 1
-                else:
-                    fps = frames / total_duration * 1000
+                    total_duration += int.from_bytes(frame_duration, "little")
+                    frames += 1
+        
+        if frames == 0:
+            return 1, 1
+        else:
+            return frames, total_duration
+    
+    @staticmethod
+    def _get_file_frames_duration_av(file: str, frames_to_iterate: Optional[int] = None, frames_only: bool = False) -> tuple[int, float]:
+        import av # type: ignore
+
+        # Getting fps from metadata is not reliable
+        # Example: https://github.com/laggykiller/sticker-convert/issues/114
+
+        with av.open(file) as container:
+            stream = container.streams.video[0]
+
+            if frames_only == True and stream.frames > 1:
+                return stream.frames, 1
+
+            last_frame = None
+            for frame_count, frame in enumerate(container.decode(stream)):
+                last_frame = frame
+
+                if frames_to_iterate != None and frame_count > frames_to_iterate:
+                    break
+
+            if frame_count <= 1:
+                return 1, 1
             else:
-                # Getting fps from metadata is not reliable
-                # Example: https://github.com/laggykiller/sticker-convert/issues/114
-                metadata = iio.immeta(file, plugin='pyav', exclude_applied=False)
-                context = None
-                if metadata.get('video_format') == 'yuv420p':
-                    if metadata.get('codec') == 'vp8':
-                        context = CodecContext.create('vp8', 'r')
-                    elif metadata.get('codec') == 'vp9':
-                        context = CodecContext.create('libvpx-vp9', 'r')
-
-                with av.open(file) as container:
-                    stream = container.streams.video[0]
-                    if not context:
-                        context = stream.codec_context
-                    
-                    frames = [i for i in container.decode(stream)]
-                    if len(frames) == 1:
-                        fps = 1
-                    else:
-                        # Need to minus one as dts is the start time of frame
-                        # Last frame dts is not last second of frame / whole video
-                        fps = (len(frames) - 1) / (frames[-1].dts * frames[-1].time_base.numerator / frames[-1].time_base.denominator)
-
-        return fps
+                duration = last_frame.pts * last_frame.time_base.numerator / last_frame.time_base.denominator * 1000
+                return frame_count, duration
 
     @staticmethod
     def get_file_codec(file: str) -> Optional[str]:
         codec = None
         try:
-            im = Image.open(file)
-            codec = im.format
+            with Image.open(file) as im:
+                codec = im.format
+                animated = im.is_animated
         except UnidentifiedImageError:
             pass
-        
-        # Unable to distinguish apng and png
-        if codec == None:
-            metadata = iio.immeta(file, plugin="pyav", exclude_applied=False)
-            codec = metadata.get("codec", None)
-            if codec == None:
-                raise RuntimeError(f"Unable to get codec for file {file}")
-            return codec
-        elif codec == "PNG":
-            if im.is_animated:
+
+        if codec == "PNG":
+            # Unable to distinguish apng and png
+            if animated:
                 return "apng"
             else:
                 return "png"
-        else:
+        elif codec != None:
             return codec.lower()
+        
+        import av # type: ignore
+        
+        with av.open(file) as container:
+            codec = container.streams.video[0].codec_context.name
+        if codec == None:
+            raise RuntimeError(f"Unable to get codec for file {file}")
+        return codec.lower()
 
     @staticmethod
     def get_file_res(file: str) -> tuple[int, int]:
         file_ext = CodecInfo.get_file_ext(file)
 
         if file_ext == ".tgs":
-            with LottieAnimation.from_tgs(file) as anim:
-                width, height = anim.lottie_animation_get_size()
-        else:
-            if file_ext in (".webp", ".png", ".apng"):
-                plugin = "pillow"
-            else:
-                plugin = "pyav"
-            frame = iio.imread(file, plugin=plugin, index=0)
-            width = frame.shape[0]
-            height = frame.shape[1]
-
-        return width, height
+            from rlottie_python import LottieAnimation  # type: ignore
 
-    @staticmethod
-    def get_file_frames(file: str) -> int:
-        file_ext = CodecInfo.get_file_ext(file)
-
-        frames = None
-
-        if file_ext == ".tgs":
             with LottieAnimation.from_tgs(file) as anim:
-                frames = anim.lottie_animation_get_totalframe()
+                width, height = anim.lottie_animation_get_size()
+        elif file_ext in (".webp", ".png", ".apng"):
+            with Image.open(file) as im:
+                width = im.width
+                height = im.height
         else:
-            if file_ext in (".webp", ".png", ".apng"):
-                frames = Image.open(file).n_frames
-            else:
-                frames = frames = len([* iio.imiter(file, plugin="pyav")])
+            import av # type: ignore
 
-        return frames
+            with av.open(file) as container:
+                stream = container.streams.video[0]
+                width = stream.width
+                height = stream.height
 
-    @staticmethod
-    def get_file_duration(file: str) -> float:
-        # Return duration in miliseconds
-        return CodecInfo.get_file_frames(file) / CodecInfo.get_file_fps(file) * 1000
+        return width, height
 
     @staticmethod
     def get_file_ext(file: str) -> str:
@@ -163,7 +229,7 @@ class CodecInfo:
 
     @staticmethod
     def is_anim(file: str) -> bool:
-        if CodecInfo.get_file_frames(file) > 1:
+        if CodecInfo.get_file_frames(file, check_anim=True) > 1:
             return True
         else:
             return False

sticker_convert/utils/media/format_verify.py

@@ -1,7 +1,5 @@
 #!/usr/bin/env python3
 import os
-import unicodedata
-import re
 from typing import Optional, Union
 
 from .codec_info import CodecInfo  # type: ignore
@@ -11,14 +9,18 @@ from ...job_option import CompOption # type: ignore
 class FormatVerify:
     @staticmethod
     def check_file(file: str, spec: CompOption) -> bool:
+        if FormatVerify.check_presence(file) == False:
+            return False
+
+        file_info = CodecInfo(file)
+
         return (
-            FormatVerify.check_presence(file)
-            and FormatVerify.check_file_res(file, res=spec.res, square=spec.square)
-            and FormatVerify.check_file_fps(file, fps=spec.fps)
-            and FormatVerify.check_file_size(file, size=spec.size_max)
-            and FormatVerify.check_animated(file, animated=spec.animated)
-            and FormatVerify.check_format(file, fmt=spec.format)
-            and FormatVerify.check_duration(file, duration=spec.duration)
+            FormatVerify.check_file_res(file, res=spec.res, square=spec.square, file_info=file_info)
+            and FormatVerify.check_file_fps(file, fps=spec.fps, file_info=file_info)
+            and FormatVerify.check_file_duration(file, duration=spec.duration, file_info=file_info)
+            and FormatVerify.check_file_size(file, size=spec.size_max, file_info=file_info)
+            and FormatVerify.check_animated(file, animated=spec.animated, file_info=file_info)
+            and FormatVerify.check_format(file, fmt=spec.format, file_info=file_info)
         )
 
     @staticmethod
@@ -29,9 +31,14 @@ class FormatVerify:
     def check_file_res(
         file: str,
         res: Optional[list[list[int]]] = None,
-        square: Optional[bool] = None
+        square: Optional[bool] = None,
+        file_info: Optional[CodecInfo] = None
     ) -> bool:
-        file_width, file_height = CodecInfo.get_file_res(file)
+        
+        if file_info:
+            file_width, file_height = file_info.res
+        else:
+            file_width, file_height = CodecInfo.get_file_res(file)
 
         if res:
             if res[0][0] and file_width < res[0][0]:
@@ -48,8 +55,15 @@ class FormatVerify:
         return True
 
     @staticmethod
-    def check_file_fps(file: str, fps: Optional[list[int]]) -> bool:
-        file_fps = CodecInfo.get_file_fps(file)
+    def check_file_fps(
+        file: str,fps: Optional[list[int]],
+        file_info: Optional[CodecInfo] = None
+    ) -> bool:
+        
+        if file_info:
+            file_fps = file_info.fps
+        else:
+            file_fps = CodecInfo.get_file_fps(file)
 
         if fps and fps[0] and file_fps < fps[0]:
             return False
@@ -57,11 +71,38 @@ class FormatVerify:
             return False
 
         return True
+    
+    @staticmethod
+    def check_file_duration(
+        file: str,
+        duration: Optional[list[str]] = None,
+        file_info: Optional[CodecInfo] = None
+    ) -> bool:
+        
+        if file_info:
+            file_duration = file_info.duration
+        else:
+            file_duration = CodecInfo.get_file_duration(file)
+
+        if duration and duration[0] and file_duration < duration[0]:
+            return False
+        if duration and duration[1] and file_duration > duration[1]:
+            return False
+
+        return True
 
     @staticmethod
-    def check_file_size(file: str, size: Optional[list[int]] = None) -> bool:
+    def check_file_size(
+        file: str,
+        size: Optional[list[int]] = None,
+        file_info: Optional[CodecInfo] = None
+    ) -> bool:
+        
         file_size = os.path.getsize(file)
-        file_animated = CodecInfo.is_anim(file)
+        if file_info:
+            file_animated = file_info.is_animated
+        else:
+            file_animated = CodecInfo.is_anim(file)
 
         if (
             file_animated == True
@@ -81,8 +122,18 @@ class FormatVerify:
         return True
 
     @staticmethod
-    def check_animated(file: str, animated: Optional[bool] = None) -> bool:
-        if animated != None and CodecInfo.is_anim(file) != animated:
+    def check_animated(
+        file: str,
+        animated: Optional[bool] = None,
+        file_info: Optional[CodecInfo] = None
+    ) -> bool:
+        
+        if file_info:
+            file_animated = file_info.is_animated
+        else:
+            file_animated = CodecInfo.is_anim(file)
+
+        if animated != None and file_animated != animated:
             return False
 
         return True
@@ -90,8 +141,17 @@ class FormatVerify:
     @staticmethod
     def check_format(
         file: str,
-        fmt: list[Union[list[str], str, None]] = None
+        fmt: list[Union[list[str], str, None]] = None,
+        file_info: Optional[CodecInfo] = None
     ):
+        
+        if file_info:
+            file_animated = file_info.is_animated
+            file_ext = file_info.file_ext
+        else:
+            file_animated = CodecInfo.is_anim(file)
+            file_ext = CodecInfo.get_file_ext(file)
+
         compat_ext = {
             ".jpg": ".jpeg",
             ".jpeg": ".jpg",
@@ -104,7 +164,7 @@ class FormatVerify:
         if fmt == [None, None]:
             return True
         
-        if FormatVerify.check_animated(file):
+        if file_animated:
             if isinstance(fmt[1], list):
                 formats = fmt[1].copy()
             else:
@@ -119,66 +179,7 @@ class FormatVerify:
             if f in compat_ext:
                 formats.append(compat_ext.get(f))  # type: ignore[arg-type]
 
-        if CodecInfo.get_file_ext(file) not in formats:
+        if file_ext not in formats:
             return False
 
         return True
-
-    @staticmethod
-    def check_duration(file: str, duration: Optional[list[str]] = None) -> bool:
-        file_duration = CodecInfo.get_file_duration(file)
-        if duration and duration[0] and file_duration < duration[0]:
-            return False
-        if duration and duration[1] and file_duration > duration[1]:
-            return False
-
-        return True
-
-    @staticmethod
-    def sanitize_filename(filename: str) -> str:
-        # Based on https://gitlab.com/jplusplus/sanitize-filename/-/blob/master/sanitize_filename/sanitize_filename.py
-        # Replace illegal character with '_'
-        """Return a fairly safe version of the filename.
-
-        We don't limit ourselves to ascii, because we want to keep municipality
-        names, etc, but we do want to get rid of anything potentially harmful,
-        and make sure we do not exceed Windows filename length limits.
-        Hence a less safe blacklist, rather than a whitelist.
-        """
-        blacklist = ["\\", "/", ":", "*", "?", '"', "<", ">", "|", "\0"]
-        reserved = [
-            "CON", "PRN", "AUX", "NUL", "COM1", "COM2", "COM3", "COM4", "COM5",
-            "COM6", "COM7", "COM8", "COM9", "LPT1", "LPT2", "LPT3", "LPT4", "LPT5",
-            "LPT6", "LPT7", "LPT8", "LPT9",
-        ]  # Reserved words on Windows
-        filename = "".join(c if c not in blacklist else "_" for c in filename)
-        # Remove all charcters below code point 32
-        filename = "".join(c if 31 < ord(c) else "_" for c in filename)
-        filename = unicodedata.normalize("NFKD", filename)
-        filename = filename.rstrip(". ")  # Windows does not allow these at end
-        filename = filename.strip()
-        if all([x == "." for x in filename]):
-            filename = "__" + filename
-        if filename in reserved:
-            filename = "__" + filename
-        if len(filename) == 0:
-            filename = "__"
-        if len(filename) > 255:
-            parts = re.split(r"/|\\", filename)[-1].split(".")
-            if len(parts) > 1:
-                ext = "." + parts.pop()
-                filename = filename[: -len(ext)]
-            else:
-                ext = ""
-            if filename == "":
-                filename = "__"
-            if len(ext) > 254:
-                ext = ext[254:]
-            maxl = 255 - len(ext)
-            filename = filename[:maxl]
-            filename = filename + ext
-            # Re-check last character (if there was no extension)
-            filename = filename.rstrip(". ")
-            if len(filename) == 0:
-                filename = "__"
-        return filename