Fix formats sorting; simplify m3u8 extraction in downloader; clean code

This commit is contained in:
Mozi 2024-10-19 17:59:18 +00:00
parent 52d9594ea6
commit 06bd726ab3
2 changed files with 39 additions and 45 deletions

View file

@ -15,7 +15,6 @@
RetryManager, RetryManager,
str_or_none, str_or_none,
traverse_obj, traverse_obj,
try_get,
urljoin, urljoin,
) )
@ -78,6 +77,7 @@ def _ws_context(self, info_dict):
""" Hold a WebSocket object and release it when leaving """ """ Hold a WebSocket object and release it when leaving """
video_id = info_dict['id'] video_id = info_dict['id']
format_id = info_dict['format_id']
live_latency = info_dict['downloader_options']['live_latency'] live_latency = info_dict['downloader_options']['live_latency']
ws_url = info_dict['downloader_options']['ws_url'] ws_url = info_dict['downloader_options']['ws_url']
@ -89,12 +89,12 @@ def _ws_context(self, info_dict):
def communicate_ws(): def communicate_ws():
self.ws = self.ydl.urlopen(Request(ws_url, headers=info_dict.get('http_headers'))) self.ws = self.ydl.urlopen(Request(ws_url, headers=info_dict.get('http_headers')))
if self.ydl.params.get('verbose', False): 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({ self.ws.send(json.dumps({
'type': 'startWatching', 'type': 'startWatching',
'data': { 'data': {
'stream': { 'stream': {
'quality': 'abr', 'quality': format_id,
'protocol': 'hls', 'protocol': 'hls',
'latency': live_latency, 'latency': live_latency,
'chasePlay': False, 'chasePlay': False,
@ -103,7 +103,6 @@ def communicate_ws():
'protocol': 'webSocket', 'protocol': 'webSocket',
'commentable': True, 'commentable': True,
}, },
'reconnect': True,
}, },
})) }))
with self.ws: with self.ws:
@ -112,7 +111,7 @@ def communicate_ws():
if not recv: if not recv:
continue continue
data = json.loads(recv) data = json.loads(recv)
if not data or not isinstance(data, dict): if not isinstance(data, dict):
continue continue
if data.get('type') == 'ping': if data.get('type') == 'ping':
# pong back # pong back
@ -126,12 +125,12 @@ def communicate_ws():
return return
elif data.get('type') == 'error': elif data.get('type') == 'error':
self.write_debug(data) 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) raise DownloadError(message)
elif self.ydl.params.get('verbose', False): elif self.ydl.params.get('verbose', False):
if len(recv) > 100: if len(recv) > 100:
recv = recv[:100] + '...' recv = recv[:100] + '...'
self.to_screen(f'[debug] Server said: {recv}') self.write_debug(f'Server said: {recv}')
stopped = threading.Event() stopped = threading.Event()
@ -146,7 +145,8 @@ def ws_main():
self.m3u8_lock.clear() # m3u8 url may be changed 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) time.sleep(self._WEBSOCKET_RECONNECT_DELAY)
self.m3u8_lock.set() # Release possible locks self.m3u8_lock.set() # Release possible locks
@ -181,7 +181,6 @@ def real_download(self, filename, info_dict):
ie = NiconicoIE(self.ydl) ie = NiconicoIE(self.ydl)
video_id = info_dict['id'] 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 # Get video info
total_duration = 0 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) retry_manager = RetryManager(self.params.get('fragment_retries'), self.report_retry)
for retry in retry_manager: for retry in retry_manager:
try: try:
# Refresh master m3u8 (if possible) and get the url of the previously-chose format # Refresh master m3u8 (if possible) to get the new URL of the previously-chose format
master_m3u8_url = ws_context._master_m3u8_url() media_m3u8_url = ie._extract_m3u8_formats(
formats = ie._extract_m3u8_formats( ws_context._master_m3u8_url(), video_id, note=False,
master_m3u8_url, video_id, query={'start': downloaded_duration}, live=False, note=False, fatal=False) query={'start': downloaded_duration}, live=False)[0]['url']
media_m3u8_url = traverse_obj(formats, (format_index, {dict}, 'url'), get_all=False)
if not media_m3u8_url:
raise DownloadError('Unable to get playlist')
# Get all fragments # Get all fragments
media_m3u8 = ie._download_webpage( media_m3u8 = ie._download_webpage(

View file

@ -7,7 +7,6 @@
import urllib.parse import urllib.parse
from .common import InfoExtractor, SearchInfoExtractor from .common import InfoExtractor, SearchInfoExtractor
from ..networking import Request
from ..networking.exceptions import HTTPError from ..networking.exceptions import HTTPError
from ..utils import ( from ..utils import (
ExtractorError, ExtractorError,
@ -957,7 +956,7 @@ class NiconicoLiveIE(NiconicoBaseIE):
def _yield_formats(self, ws_url, headers, latency, video_id, is_live): def _yield_formats(self, ws_url, headers, latency, video_id, is_live):
ws = self._request_webpage( 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') self.write_debug('Sending HLS server request')
ws.send(json.dumps({ ws.send(json.dumps({
@ -973,37 +972,36 @@ def _yield_formats(self, ws_url, headers, latency, video_id, is_live):
'protocol': 'webSocket', 'protocol': 'webSocket',
'commentable': True, 'commentable': True,
}, },
'reconnect': False,
}, },
})) }))
while True: with ws:
recv = ws.recv() while True:
if not recv: recv = ws.recv()
continue if not recv:
data = json.loads(recv) continue
if not isinstance(data, dict): data = json.loads(recv)
continue if not isinstance(data, dict):
if data.get('type') == 'stream': continue
m3u8_url = data['data']['uri'] if data.get('type') == 'stream':
qualities = data['data']['availableQualities'] m3u8_url = data['data']['uri']
break qualities = data['data']['availableQualities']
elif data.get('type') == 'disconnect': break
self.write_debug(recv) elif data.get('type') == 'disconnect':
raise ExtractorError('Disconnected at middle of extraction') self.write_debug(data)
elif data.get('type') == 'error': raise ExtractorError('Disconnected at middle of extraction')
self.write_debug(recv) elif data.get('type') == 'error':
message = traverse_obj(data, ('body', 'code')) or recv self.write_debug(data)
raise ExtractorError(message) message = traverse_obj(data, ('data', 'code')) or recv
elif self.get_param('verbose', False): raise ExtractorError(message)
if len(recv) > 100: elif self.get_param('verbose', False):
recv = recv[:100] + '...' if len(recv) > 100:
self.write_debug(f'Server said: {recv}') recv = recv[:100] + '...'
self.write_debug(f'Server said: {recv}')
ws.close() formats = sorted(self._extract_m3u8_formats(
m3u8_url, video_id, ext='mp4', live=is_live), key=lambda f: f['tbr'], reverse=True)
formats = self._extract_m3u8_formats(m3u8_url, video_id, ext='mp4', live=is_live) for fmt, q in zip(formats, qualities[1:]):
for fmt, q in zip(formats, reversed(qualities[1:])):
fmt.update({ fmt.update({
'format_id': q, 'format_id': q,
'protocol': 'niconico_live', 'protocol': 'niconico_live',