[mhtml, cleanup] Use imghdr

[yt-dlp.git] / yt_dlp / downloader / mhtml.py

diff --git a/yt_dlp/downloader/mhtml.py b/yt_dlp/downloader/mhtml.py
index ce2d399470540c8f4467b9e7c71d3a14ae0c303b..ed076e09ed7d3cec75a66a5146089c2660b8058b 100644 (file)
--- a/yt_dlp/downloader/mhtml.py
+++ b/yt_dlp/downloader/mhtml.py
@@ -4,6 +4,7 @@
 import uuid
 
 from .fragment import FragmentFD
+from ..compat import imghdr
 from ..utils import escapeHTML, formatSeconds, srt_subtitles_timecode, urljoin
 from ..version import __version__ as YT_DLP_VERSION
 
@@ -166,21 +167,13 @@ def real_download(self, filename, info_dict):
                 continue
             frag_content = self._read_fragment(ctx)
 
-            mime_type = b'image/jpeg'
-            if frag_content.startswith(b'\x89PNG\r\n\x1a\n'):
-                mime_type = b'image/png'
-            if frag_content.startswith((b'GIF87a', b'GIF89a')):
-                mime_type = b'image/gif'
-            if frag_content.startswith(b'RIFF') and frag_content[8:12] == b'WEBP':
-                mime_type = b'image/webp'
-
             frag_header = io.BytesIO()
             frag_header.write(
                 b'--%b\r\n' % frag_boundary.encode('us-ascii'))
             frag_header.write(
                 b'Content-ID: <%b>\r\n' % self._gen_cid(i, fragment, frag_boundary).encode('us-ascii'))
             frag_header.write(
-                b'Content-type: %b\r\n' % mime_type)
+                b'Content-type: %b\r\n' % f'image/{imghdr.what(h=frag_content) or "jpeg"}'.encode())
             frag_header.write(
                 b'Content-length: %u\r\n' % len(frag_content))
             frag_header.write(