From f2a4ea1794718e4dc0148bc172cb877f1080903b Mon Sep 17 00:00:00 2001 From: Simon Sawicki Date: Thu, 27 Jun 2024 16:12:19 +0200 Subject: [PATCH] [pp/embedthumbnail] Fix postprocessor (#10248) * [compat] Improve `imghdr.what` detection * [pp/embedthumbnail] Improve imghdr fail message * [pp/embedthumbnail] Fix AtomicParsley error handling Authored by: Grub4K --- yt_dlp/compat/imghdr.py | 26 ++++++++++++++++---------- yt_dlp/postprocessor/embedthumbnail.py | 18 +++++++++++++----- 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/yt_dlp/compat/imghdr.py b/yt_dlp/compat/imghdr.py index 5d64ab07b..4ae173fde 100644 --- a/yt_dlp/compat/imghdr.py +++ b/yt_dlp/compat/imghdr.py @@ -1,16 +1,22 @@ -tests = { - 'webp': lambda h: h[0:4] == b'RIFF' and h[8:] == b'WEBP', - 'png': lambda h: h[:8] == b'\211PNG\r\n\032\n', - 'jpeg': lambda h: h[6:10] in (b'JFIF', b'Exif'), - 'gif': lambda h: h[:6] in (b'GIF87a', b'GIF89a'), -} - - def what(file=None, h=None): """Detect format of image (Currently supports jpeg, png, webp, gif only) - Ref: https://github.com/python/cpython/blob/3.10/Lib/imghdr.py + Ref: https://github.com/python/cpython/blob/3.11/Lib/imghdr.py + Ref: https://www.w3.org/Graphics/JPEG/itu-t81.pdf """ if h is None: with open(file, 'rb') as f: h = f.read(12) - return next((type_ for type_, test in tests.items() if test(h)), None) + + if h.startswith(b'RIFF') and h.startswith(b'WEBP', 8): + return 'webp' + + if h.startswith(b'\x89PNG'): + return 'png' + + if h.startswith(b'\xFF\xD8\xFF'): + return 'jpeg' + + if h.startswith(b'GIF'): + return 'gif' + + return None diff --git a/yt_dlp/postprocessor/embedthumbnail.py b/yt_dlp/postprocessor/embedthumbnail.py index 673a92468..f2228ac61 100644 --- a/yt_dlp/postprocessor/embedthumbnail.py +++ b/yt_dlp/postprocessor/embedthumbnail.py @@ -119,14 +119,21 @@ def run(self, info): if not mutagen or prefer_atomicparsley: success = False else: + self._report_run('mutagen', filename) + f = {'jpeg': MP4Cover.FORMAT_JPEG, 'png': MP4Cover.FORMAT_PNG} try: - self._report_run('mutagen', filename) + with open(thumbnail_filename, 'rb') as thumbfile: + thumb_data = thumbfile.read() + + type_ = imghdr.what(h=thumb_data) + if not type_: + raise ValueError('could not determine image type') + elif type_ not in f: + raise ValueError(f'incompatible image type: {type_}') + meta = MP4(filename) # NOTE: the 'covr' atom is a non-standard MPEG-4 atom, # Apple iTunes 'M4A' files include the 'moov.udta.meta.ilst' atom. - f = {'jpeg': MP4Cover.FORMAT_JPEG, 'png': MP4Cover.FORMAT_PNG}[imghdr.what(thumbnail_filename)] - with open(thumbnail_filename, 'rb') as thumbfile: - thumb_data = thumbfile.read() meta.tags['covr'] = [MP4Cover(data=thumb_data, imageformat=f)] meta.save() temp_filename = filename @@ -160,9 +167,10 @@ def run(self, info): stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if returncode: self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {stderr.strip()}') + success = False # for formats that don't support thumbnails (like 3gp) AtomicParsley # won't create to the temporary file - if 'No changes' in stdout: + elif 'No changes' in stdout: self.report_warning('The file format doesn\'t support embedding a thumbnail') success = False