[niconico:live] connect to WS endpoint for better format extraction

yt_dlp/downloader/niconico.py

@@ -11,7 +11,6 @@
 from ..utils import (
     str_or_none,
     std_headers,
-    to_str,
     DownloadError,
     try_get,
 )
@@ -78,29 +77,27 @@ def real_download(self, filename, info_dict):
 
         video_id = info_dict['video_id']
         ws_url = info_dict['url']
+        ws_extractor = info_dict['ws']
         cookies = info_dict.get('cookies')
         live_quality = info_dict.get('live_quality', 'high')
         live_latency = info_dict.get('live_latency', 'high')
         dl = FFmpegFD(self.ydl, self.params or {})
-        self.to_screen('[%s] %s: Fetching HLS playlist info via WebSocket' % ('niconico:live', video_id))
 
         new_info_dict = info_dict.copy()
         new_info_dict.update({
-            'url': None,
             'protocol': 'live_ffmpeg',
         })
-        lock = threading.Lock()
-        lock.acquire()
 
         def communicate_ws(reconnect):
-            with WebSocket(ws_url, {
-                'Cookie': str_or_none(cookies) or '',
-                'Origin': 'https://live2.nicovideo.jp',
-                'Accept': '*/*',
-                'User-Agent': std_headers['User-Agent'],
-            }) as ws:
+            if reconnect:
+                ws = WebSocket(ws_url, {
+                    'Cookie': str_or_none(cookies) or '',
+                    'Origin': 'https://live2.nicovideo.jp',
+                    'Accept': '*/*',
+                    'User-Agent': std_headers['User-Agent'],
+                })
                 if self.ydl.params.get('verbose', False):
-                    self.to_screen('[debug] Sending HLS server request')
+                    self.to_screen('[debug] Sending startWatching request')
                 ws.send(json.dumps({
                     "type": "startWatching",
                     "data": {
@@ -114,29 +111,28 @@ def communicate_ws(reconnect):
                             "protocol": "webSocket",
                             "commentable": True
                         },
-                        "reconnect": reconnect,
+                        "reconnect": True,
                     }
                 }))
-
+            else:
+                ws = ws_extractor
+            with ws:
                 while True:
-                    recv = to_str(ws.recv()).strip()
+                    recv = ws.recv()
                     if not recv:
                         continue
                     data = json.loads(recv)
                     if not data or not isinstance(data, dict):
                         continue
-                    if data.get('type') == 'stream' and not new_info_dict.get('url'):
-                        new_info_dict['url'] = data['data']['uri']
-                        lock.release()
-                    elif data.get('type') == 'ping':
+                    if data.get('type') == 'ping':
                         # pong back
                         ws.send(r'{"type":"pong"}')
                         ws.send(r'{"type":"keepSeat"}')
                     elif data.get('type') == 'disconnect':
-                        print(data)
+                        self.write_debug(data)
                         return True
                     elif data.get('type') == 'error':
-                        print(data)
+                        self.write_debug(data)
                         message = try_get(data, lambda x: x["body"]["code"], compat_str) or recv
                         return DownloadError(message)
                     elif self.ydl.params.get('verbose', False):
@@ -151,10 +147,6 @@ def ws_main():
                     ret = communicate_ws(reconnect)
                     if ret is True:
                         return
-                    if isinstance(ret, BaseException):
-                        new_info_dict['error'] = ret
-                        lock.release()
-                        return
                 except BaseException as e:
                     self.to_screen('[%s] %s: Connection error occured, reconnecting after 10 seconds: %s' % ('niconico:live', video_id, str_or_none(e)))
                     time.sleep(10)
@@ -165,8 +157,4 @@ def ws_main():
         thread = threading.Thread(target=ws_main, daemon=True)
         thread.start()
 
-        lock.acquire(True)
-        err = new_info_dict.get('error')
-        if isinstance(err, BaseException):
-            raise err
         return dl.download(filename, new_info_dict)

yt_dlp/extractor/niconico.py

@@ -27,13 +27,16 @@
     parse_duration,
     parse_iso8601,
     remove_start,
-    try_get,
+    std_headers,
+    str_or_none,
+    time_millis,
     traverse_obj,
+    try_get,
     unescapeHTML,
-    urlencode_postdata,
     update_url_query,
-    time_millis,
+    urlencode_postdata,
 )
+from ..websocket import WebSocket
 
 
 class NiconicoBaseIE(InfoExtractor):
@@ -984,16 +987,6 @@ class NiconicoLiveIE(NiconicoBaseIE):
     _VALID_URL = r'(?:https?://(?:sp\.)?live2?\.nicovideo\.jp/(?:watch|gate)/|nico(?:nico|video)?:)(?P<id>lv\d+)'
     _FEATURE_DEPENDENCY = ('websocket', )
 
-    # sort qualities in this order to trick youtube-dl to download highest quality as default
-    _KNOWN_QUALITIES = (
-        # quality code, group (bitmask), tags...
-        ('abr', 3, 'all'),
-        ('super_low', 1, 'all'),
-        ('low', 1, 'all'),
-        ('normal', 1, 'all'),
-        ('high', 1, 'all'),
-        ('super_high', 1, 'all'),
-    )
     _KNOWN_LATENCY = ('high', 'low')
 
     def _real_extract(self, url):
@@ -1011,31 +1004,80 @@ def _real_extract(self, url):
             'frontend_id': '9',
         })
 
+        cookies = try_get(urlh.geturl(), self._get_cookie_header)
         latency = try_get(self._configuration_arg('latency'), lambda x: x[0])
         if latency not in self._KNOWN_LATENCY:
             latency = 'high'
 
+        ws = WebSocket(ws_url, {
+            'Cookie': str_or_none(cookies) or '',
+            'Origin': 'https://live2.nicovideo.jp',
+            'Accept': '*/*',
+            'User-Agent': std_headers['User-Agent'],
+        })
+
+        self.write_debug('[debug] Sending HLS server request')
+        ws.send(json.dumps({
+            "type": "startWatching",
+            "data": {
+                "stream": {
+                    "quality": 'abr',
+                    "protocol": "hls+fmp4",
+                    "latency": latency,
+                    "chasePlay": False
+                },
+                "room": {
+                    "protocol": "webSocket",
+                    "commentable": True
+                },
+                "reconnect": False,
+            }
+        }))
+
+        while True:
+            recv = ws.recv()
+            if not recv:
+                continue
+            data = json.loads(recv)
+            if not data or not isinstance(data, dict):
+                continue
+            if data.get('type') == 'stream':
+                m3u8_url = data['data']['uri']
+                qualities = data['data']['availableQualities']
+                break
+            elif data.get('type') == 'disconnect':
+                self.write_debug(recv)
+                raise ExtractorError('Disconnected at middle of extraction')
+            elif data.get('type') == 'error':
+                self.write_debug(recv)
+                message = try_get(data, lambda x: x["body"]["code"], compat_str) or recv
+                raise ExtractorError(message)
+            elif self.get_param('verbose', False):
+                if len(recv) > 100:
+                    recv = recv[:100] + '...'
+                self.to_screen('[debug] Server said: %s' % recv)
+
         title = try_get(
             None,
             (lambda x: embedded_data['program']['title'],
              lambda x: self._html_search_meta(('og:title', 'twitter:title'), webpage, 'live title', fatal=False)),
             compat_str)
 
-        cookies = try_get(urlh.geturl(), self._get_cookie_header)
-
-        return {
-            'id': video_id,
-            'title': title,
-            'formats': [{
-                'url': ws_url,
+        formats = self._extract_m3u8_formats(m3u8_url, video_id, ext='mp4', live=True)
+        self._sort_formats(formats)
+        for fmt, q in zip(formats, reversed(qualities[1:])):
+            fmt.update({
+                'format_id': q,
                 'protocol': 'niconico_live',
-                'format_id': '%s' % quality[0],
+                'ws': ws,
                 'video_id': video_id,
                 'cookies': cookies,
-                'ext': 'mp4',
-                'is_live': True,
-                'live_quality': quality,
                 'live_latency': latency,
-            } for quality in self._KNOWN_QUALITIES],
+            })
+
+        return {
+            'id': video_id,
+            'title': title,
+            'formats': formats,
             'is_live': True,
         }