[niconico] Implement heartbeat for download

2 files changed, 83 insertions, 8 deletions

diff --git a/youtube_dl/downloader/__init__.py b/youtube_dl/downloader/__init__.py
index 2e485df9d..d8f2fa342 100644
--- a/youtube_dl/downloader/__init__.py
+++ b/youtube_dl/downloader/__init__.py
@@ -1,22 +1,31 @@
 from __future__ import unicode_literals
 
+from ..utils import (
+    determine_protocol,
+)
+
+
+def get_suitable_downloader(info_dict, params={}):
+    info_dict['protocol'] = determine_protocol(info_dict)
+    info_copy = info_dict.copy()
+    return _get_suitable_downloader(info_copy, params)
+
+
+# Some of these require get_suitable_downloader
 from .common import FileDownloader
+from .dash import DashSegmentsFD
 from .f4m import F4mFD
 from .hls import HlsFD
 from .http import HttpFD
 from .rtmp import RtmpFD
-from .dash import DashSegmentsFD
 from .rtsp import RtspFD
 from .ism import IsmFD
+from .niconico import NiconicoDmcFD
 from .external import (
     get_external_downloader,
     FFmpegFD,
 )
 
-from ..utils import (
-    determine_protocol,
-)
-
 PROTOCOL_MAP = {
     'rtmp': RtmpFD,
     'm3u8_native': HlsFD,
@@ -26,13 +35,12 @@ PROTOCOL_MAP = {
     'f4m': F4mFD,
     'http_dash_segments': DashSegmentsFD,
     'ism': IsmFD,
+    'niconico_dmc': NiconicoDmcFD,
 }
 
 
-def get_suitable_downloader(info_dict, params={}):
+def _get_suitable_downloader(info_dict, params={}):
     """Get the downloader class that can handle the info dict."""
-    protocol = determine_protocol(info_dict)
-    info_dict['protocol'] = protocol
 
     # 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
@@ -43,6 +51,7 @@ def get_suitable_downloader(info_dict, params={}):
         if ed.can_download(info_dict):
             return ed
 
+    protocol = info_dict['protocol']
     if protocol.startswith('m3u8') and info_dict.get('is_live'):
         return FFmpegFD
 
diff --git a/youtube_dl/downloader/niconico.py b/youtube_dl/downloader/niconico.py
new file mode 100644
index 000000000..6392c9989
--- /dev/null
+++ b/youtube_dl/downloader/niconico.py
@@ -0,0 +1,66 @@
+# coding: utf-8
+from __future__ import unicode_literals
+
+try:
+    import threading
+except ImportError:
+    threading = None
+
+from .common import FileDownloader
+from ..downloader import get_suitable_downloader
+from ..extractor.niconico import NiconicoIE
+from ..utils import sanitized_Request
+
+
+class NiconicoDmcFD(FileDownloader):
+    """ Downloading niconico douga from DMC with heartbeat """
+
+    FD_NAME = 'niconico_dmc'
+
+    def real_download(self, filename, info_dict):
+        self.to_screen('[%s] Downloading from DMC' % self.FD_NAME)
+
+        ie = NiconicoIE(self.ydl)
+        info_dict, heartbeat_info_dict = ie._get_heartbeat_info(info_dict)
+
+        fd = get_suitable_downloader(info_dict, params=self.params)(self.ydl, self.params)
+        for ph in self._progress_hooks:
+            fd.add_progress_hook(ph)
+
+        if not threading:
+            self.to_screen('[%s] Threading for Heartbeat not available' % self.FD_NAME)
+            return fd.real_download(filename, info_dict)
+
+        success = download_complete = False
+        timer = [None]
+        heartbeat_lock = threading.Lock()
+        heartbeat_url = heartbeat_info_dict['url']
+        heartbeat_data = heartbeat_info_dict['data'].encode()
+        heartbeat_interval = heartbeat_info_dict.get('interval', 30)
+
+        request = sanitized_Request(heartbeat_url, heartbeat_data)
+
+        def heartbeat():
+            try:
+                self.ydl.urlopen(request).read()
+            except Exception:
+                self.to_screen('[%s] Heartbeat failed' % self.FD_NAME)
+
+            with heartbeat_lock:
+                if not download_complete:
+                    timer[0] = threading.Timer(heartbeat_interval, heartbeat)
+                    timer[0].start()
+
+        heartbeat_info_dict['ping']()
+        self.to_screen('[%s] Heartbeat with %d second interval ...' % (self.FD_NAME, heartbeat_interval))
+        try:
+            heartbeat()
+            if type(fd).__name__ == 'HlsFD':
+                info_dict.update(ie._extract_m3u8_formats(info_dict['url'], info_dict['id'])[0])
+            success = fd.real_download(filename, info_dict)
+        finally:
+            if heartbeat_lock:
+                with heartbeat_lock:
+                    timer[0].cancel()
+                    download_complete = True
+            return success