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,
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(

View file

@ -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',