From 06bd726ab312008d3d54f5ac214c0180975e0ad6 Mon Sep 17 00:00:00 2001 From: Mozi <29089388+pzhlkj6612@users.noreply.github.com> Date: Sat, 19 Oct 2024 17:59:18 +0000 Subject: [PATCH] Fix formats sorting; simplify m3u8 extraction in downloader; clean code --- yt_dlp/downloader/niconico.py | 28 ++++++++---------- yt_dlp/extractor/niconico.py | 56 +++++++++++++++++------------------ 2 files changed, 39 insertions(+), 45 deletions(-) diff --git a/yt_dlp/downloader/niconico.py b/yt_dlp/downloader/niconico.py index 17a8f8f91b..e8b0d53423 100644 --- a/yt_dlp/downloader/niconico.py +++ b/yt_dlp/downloader/niconico.py @@ -15,7 +15,6 @@ RetryManager, str_or_none, traverse_obj, - try_get, urljoin, ) @@ -78,6 +77,7 @@ def _ws_context(self, info_dict): """ Hold a WebSocket object and release it when leaving """ video_id = info_dict['id'] + format_id = info_dict['format_id'] live_latency = info_dict['downloader_options']['live_latency'] ws_url = info_dict['downloader_options']['ws_url'] @@ -89,12 +89,12 @@ def _ws_context(self, info_dict): def communicate_ws(): self.ws = self.ydl.urlopen(Request(ws_url, headers=info_dict.get('http_headers'))) if self.ydl.params.get('verbose', False): - self.to_screen('[debug] Sending startWatching request') + self.write_debug('Sending HLS server request') self.ws.send(json.dumps({ 'type': 'startWatching', 'data': { 'stream': { - 'quality': 'abr', + 'quality': format_id, 'protocol': 'hls', 'latency': live_latency, 'chasePlay': False, @@ -103,7 +103,6 @@ def communicate_ws(): 'protocol': 'webSocket', 'commentable': True, }, - 'reconnect': True, }, })) with self.ws: @@ -112,7 +111,7 @@ def communicate_ws(): if not recv: continue data = json.loads(recv) - if not data or not isinstance(data, dict): + if not isinstance(data, dict): continue if data.get('type') == 'ping': # pong back @@ -126,12 +125,12 @@ def communicate_ws(): return elif data.get('type') == 'error': self.write_debug(data) - message = try_get(data, lambda x: x['body']['code'], str) or recv + message = traverse_obj(data, ('data', 'code')) or recv raise DownloadError(message) elif self.ydl.params.get('verbose', False): if len(recv) > 100: recv = recv[:100] + '...' - self.to_screen(f'[debug] Server said: {recv}') + self.write_debug(f'Server said: {recv}') stopped = threading.Event() @@ -146,7 +145,8 @@ def ws_main(): self.m3u8_lock.clear() # m3u8 url may be changed - self.to_screen('[{}] {}: Connection error occured, reconnecting after {} seconds: {}'.format('niconico:live', video_id, self._WEBSOCKET_RECONNECT_DELAY, str_or_none(e))) + self.to_screen('[{}] {}: Connection error occured, reconnecting after {} seconds: {}'.format( + 'niconico:live', video_id, self._WEBSOCKET_RECONNECT_DELAY, str_or_none(e))) time.sleep(self._WEBSOCKET_RECONNECT_DELAY) self.m3u8_lock.set() # Release possible locks @@ -181,7 +181,6 @@ def real_download(self, filename, info_dict): ie = NiconicoIE(self.ydl) video_id = info_dict['id'] - format_index = next((i for i, fmt in enumerate(info_dict['formats']) if fmt['format_id'] == info_dict['format_id'])) # Get video info total_duration = 0 @@ -209,13 +208,10 @@ def real_download(self, filename, info_dict): retry_manager = RetryManager(self.params.get('fragment_retries'), self.report_retry) for retry in retry_manager: try: - # Refresh master m3u8 (if possible) and get the url of the previously-chose format - master_m3u8_url = ws_context._master_m3u8_url() - formats = ie._extract_m3u8_formats( - master_m3u8_url, video_id, query={'start': downloaded_duration}, live=False, note=False, fatal=False) - media_m3u8_url = traverse_obj(formats, (format_index, {dict}, 'url'), get_all=False) - if not media_m3u8_url: - raise DownloadError('Unable to get playlist') + # Refresh master m3u8 (if possible) to get the new URL of the previously-chose format + media_m3u8_url = ie._extract_m3u8_formats( + ws_context._master_m3u8_url(), video_id, note=False, + query={'start': downloaded_duration}, live=False)[0]['url'] # Get all fragments media_m3u8 = ie._download_webpage( diff --git a/yt_dlp/extractor/niconico.py b/yt_dlp/extractor/niconico.py index f39d4d3ba6..61beb4515d 100644 --- a/yt_dlp/extractor/niconico.py +++ b/yt_dlp/extractor/niconico.py @@ -7,7 +7,6 @@ import urllib.parse from .common import InfoExtractor, SearchInfoExtractor -from ..networking import Request from ..networking.exceptions import HTTPError from ..utils import ( ExtractorError, @@ -957,7 +956,7 @@ class NiconicoLiveIE(NiconicoBaseIE): def _yield_formats(self, ws_url, headers, latency, video_id, is_live): ws = self._request_webpage( - Request(ws_url, headers=headers), video_id, note='Connecting to WebSocket server') + ws_url, video_id, note='Connecting to WebSocket server', headers=headers) self.write_debug('Sending HLS server request') ws.send(json.dumps({ @@ -973,37 +972,36 @@ def _yield_formats(self, ws_url, headers, latency, video_id, is_live): 'protocol': 'webSocket', 'commentable': True, }, - 'reconnect': False, }, })) - while True: - recv = ws.recv() - if not recv: - continue - data = json.loads(recv) - if 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 = traverse_obj(data, ('body', 'code')) or recv - raise ExtractorError(message) - elif self.get_param('verbose', False): - if len(recv) > 100: - recv = recv[:100] + '...' - self.write_debug(f'Server said: {recv}') + with ws: + while True: + recv = ws.recv() + if not recv: + continue + data = json.loads(recv) + if 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(data) + raise ExtractorError('Disconnected at middle of extraction') + elif data.get('type') == 'error': + self.write_debug(data) + message = traverse_obj(data, ('data', 'code')) or recv + raise ExtractorError(message) + elif self.get_param('verbose', False): + if len(recv) > 100: + recv = recv[:100] + '...' + self.write_debug(f'Server said: {recv}') - ws.close() - - formats = self._extract_m3u8_formats(m3u8_url, video_id, ext='mp4', live=is_live) - for fmt, q in zip(formats, reversed(qualities[1:])): + formats = sorted(self._extract_m3u8_formats( + m3u8_url, video_id, ext='mp4', live=is_live), key=lambda f: f['tbr'], reverse=True) + for fmt, q in zip(formats, qualities[1:]): fmt.update({ 'format_id': q, 'protocol': 'niconico_live',