yt-dlp/yt_dlp/downloader/websocket.py
2022-07-14 09:11:54 +09:00

94 lines
3.2 KiB
Python

import contextlib
import os
import threading
import time
from .common import FileDownloader
from .external import FFmpegFD
from ..compat import asyncio
from ..dependencies import websockets
class _WebSocketFD(FileDownloader):
async def connect(self, stdin, info_dict):
try:
await self.real_connection(stdin, info_dict)
except OSError:
pass
finally:
with contextlib.suppress(OSError):
stdin.flush()
stdin.close()
async def real_connection(self, sink, info_dict):
async with websockets.connect(info_dict['url'], extra_headers=info_dict.get('http_headers', {})) as ws:
while True:
recv = await ws.recv()
if isinstance(recv, str):
recv = recv.encode('utf8')
sink.write(recv)
class WebSocketFragmentFD(_WebSocketFD):
""" A sink to ffmpeg for downloading fragments in any form """
def real_download(self, filename, info_dict):
info_copy = info_dict.copy()
info_copy['url'] = '-'
connect = self.connect
class FFmpegStdinFD(FFmpegFD):
@classmethod
def get_basename(cls):
return FFmpegFD.get_basename()
def on_process_started(self, proc, stdin):
thread = threading.Thread(target=asyncio.run, daemon=True, args=(connect(stdin, info_dict), ))
thread.start()
return FFmpegStdinFD(self.ydl, self.params or {}).download(filename, info_copy)
class WebSocketToFileFD(_WebSocketFD):
""" A sink to a file for downloading fragments in any form """
def real_download(self, filename, info_dict):
tempname = self.temp_name(filename)
try:
with open(tempname, 'wb') as w:
started = time.time()
status = {
'filename': info_dict.get('_filename'),
'status': 'downloading',
'elapsed': 0,
'downloaded_bytes': 0,
}
self._hook_progress(status, info_dict)
thread = threading.Thread(target=asyncio.run, daemon=True, args=(self.connect(w, info_dict), ))
thread.start()
time_and_size, avg_len = [], 10
while thread.is_alive():
time.sleep(0.1)
downloaded, curr = w.tell(), time.time()
# taken from ffmpeg attachment
time_and_size.append((downloaded, curr))
time_and_size = time_and_size[-avg_len:]
if len(time_and_size) > 1:
last, early = time_and_size[0], time_and_size[-1]
average_speed = (early[0] - last[0]) / (early[1] - last[1])
else:
average_speed = None
status.update({
'downloaded_bytes': downloaded,
'speed': average_speed,
'elapsed': curr - started,
})
self._hook_progress(status, info_dict)
except KeyboardInterrupt:
pass
finally:
os.replace(tempname, filename)
return True