[cleanup] Minor fixes (See desc)

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

diff --git a/yt_dlp/downloader/niconico.py b/yt_dlp/downloader/niconico.py
index dc49dff585cb1e2bcd995aa90b1f2aecc06d1c0d..5947446b1409f5272b4eb489d9e1e688c5d6886c 100644 (file)
--- a/yt_dlp/downloader/niconico.py
+++ b/yt_dlp/downloader/niconico.py
@@ -1,12 +1,9 @@
-# coding: utf-8
-from __future__ import unicode_literals
-
 import threading
 
 from .common import FileDownloader
-from ..downloader import _get_real_downloader
+from ..downloader import get_suitable_downloader
 from ..extractor.niconico import NiconicoIE
-from ..compat import compat_urllib_request
+from ..utils import sanitized_Request
 
 
 class NiconicoDmcFD(FileDownloader):
@@ -20,20 +17,20 @@ def real_download(self, filename, info_dict):
         ie = NiconicoIE(self.ydl)
         info_dict, heartbeat_info_dict = ie._get_heartbeat_info(info_dict)
 
-        fd = _get_real_downloader(info_dict, params=self.params)(self.ydl, self.params)
+        fd = get_suitable_downloader(info_dict, params=self.params)(self.ydl, self.params)
 
         success = download_complete = False
         timer = [None]
-
         heartbeat_lock = threading.Lock()
         heartbeat_url = heartbeat_info_dict['url']
-        heartbeat_data = heartbeat_info_dict['data']
+        heartbeat_data = heartbeat_info_dict['data'].encode()
         heartbeat_interval = heartbeat_info_dict.get('interval', 30)
-        self.to_screen('[%s] Heartbeat with %s second interval ...' % (self.FD_NAME, heartbeat_interval))
+
+        request = sanitized_Request(heartbeat_url, heartbeat_data)
 
         def heartbeat():
             try:
-                compat_urllib_request.urlopen(url=heartbeat_url, data=heartbeat_data.encode())
+                self.ydl.urlopen(request).read()
             except Exception:
                 self.to_screen('[%s] Heartbeat failed' % self.FD_NAME)
 
@@ -42,13 +39,16 @@ def heartbeat():
                     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
+        return success