[downloader/mhtml] Add new downloader (#343)

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

diff --git a/yt_dlp/downloader/__init__.py b/yt_dlp/downloader/__init__.py
index c2e155c0ac120c7f59937fe621cca865fa7b6e06..82d7623f62656b8fabc32ac82386d8eb482fa129 100644 (file)
--- a/yt_dlp/downloader/__init__.py
+++ b/yt_dlp/downloader/__init__.py
@@ -1,5 +1,6 @@
 from __future__ import unicode_literals
 
+from ..compat import compat_str
 from ..utils import (
     determine_protocol,
 )
@@ -21,6 +22,7 @@ def _get_real_downloader(info_dict, protocol=None, *args, **kwargs):
 from .rtmp import RtmpFD
 from .rtsp import RtspFD
 from .ism import IsmFD
+from .mhtml import MhtmlFD
 from .niconico import NiconicoDmcFD
 from .youtube_live_chat import YoutubeLiveChatReplayFD
 from .external import (
@@ -30,6 +32,7 @@ def _get_real_downloader(info_dict, protocol=None, *args, **kwargs):
 
 PROTOCOL_MAP = {
     'rtmp': RtmpFD,
+    'rtmp_ffmpeg': FFmpegFD,
     'm3u8_native': HlsFD,
     'm3u8': FFmpegFD,
     'mms': RtspFD,
@@ -37,11 +40,31 @@ def _get_real_downloader(info_dict, protocol=None, *args, **kwargs):
     'f4m': F4mFD,
     'http_dash_segments': DashSegmentsFD,
     'ism': IsmFD,
+    'mhtml': MhtmlFD,
     'niconico_dmc': NiconicoDmcFD,
     'youtube_live_chat_replay': YoutubeLiveChatReplayFD,
 }
 
 
+def shorten_protocol_name(proto, simplify=False):
+    short_protocol_names = {
+        'm3u8_native': 'm3u8_n',
+        'rtmp_ffmpeg': 'rtmp_f',
+        'http_dash_segments': 'dash',
+        'niconico_dmc': 'dmc',
+    }
+    if simplify:
+        short_protocol_names.update({
+            'https': 'http',
+            'ftps': 'ftp',
+            'm3u8_native': 'm3u8',
+            'rtmp_ffmpeg': 'rtmp',
+            'm3u8_frag_urls': 'm3u8',
+            'dash_frag_urls': 'dash',
+        })
+    return short_protocol_names.get(proto, proto)
+
+
 def get_suitable_downloader(info_dict, params={}, default=HttpFD):
     """Get the downloader class that can handle the info dict."""
     protocol = determine_protocol(info_dict)
@@ -50,16 +73,24 @@ def get_suitable_downloader(info_dict, params={}, default=HttpFD):
     # if (info_dict.get('start_time') or info_dict.get('end_time')) and not info_dict.get('requested_formats') and FFmpegFD.can_download(info_dict):
     #     return FFmpegFD
 
-    external_downloader = params.get('external_downloader')
-    if external_downloader is not None:
+    downloaders = params.get('external_downloader')
+    external_downloader = (
+        downloaders if isinstance(downloaders, compat_str) or downloaders is None
+        else downloaders.get(shorten_protocol_name(protocol, True), downloaders.get('default')))
+    if external_downloader and external_downloader.lower() == 'native':
+        external_downloader = 'native'
+
+    if external_downloader not in (None, 'native'):
         ed = get_external_downloader(external_downloader)
         if ed.can_download(info_dict, external_downloader):
             return ed
 
-    if protocol.startswith('m3u8'):
+    if protocol in ('m3u8', 'm3u8_native'):
         if info_dict.get('is_live'):
             return FFmpegFD
-        elif _get_real_downloader(info_dict, 'frag_urls', params, None):
+        elif external_downloader == 'native':
+            return HlsFD
+        elif _get_real_downloader(info_dict, 'm3u8_frag_urls', params, None):
             return HlsFD
         elif params.get('hls_prefer_native') is True:
             return HlsFD
@@ -70,6 +101,7 @@ def get_suitable_downloader(info_dict, params={}, default=HttpFD):
 
 
 __all__ = [
-    'get_suitable_downloader',
     'FileDownloader',
+    'get_suitable_downloader',
+    'shorten_protocol_name',
 ]