mirror of
https://github.com/yt-dlp/yt-dlp.git
synced 2025-09-03 00:25:08 +00:00
[rh:websockets] Migrate websockets to networking framework (#7720)
* Adds a basic WebSocket framework * Introduces new minimum `websockets` version of 12.0 * Deprecates `WebSocketsWrapper` Fixes https://github.com/yt-dlp/yt-dlp/issues/8439 Authored by: coletdjnz
This commit is contained in:
@ -1,5 +1,3 @@
|
||||
import asyncio
|
||||
import atexit
|
||||
import base64
|
||||
import binascii
|
||||
import calendar
|
||||
@ -54,7 +52,7 @@ from ..compat import (
|
||||
compat_os_name,
|
||||
compat_shlex_quote,
|
||||
)
|
||||
from ..dependencies import websockets, xattr
|
||||
from ..dependencies import xattr
|
||||
|
||||
__name__ = __name__.rsplit('.', 1)[0] # Pretend to be the parent module
|
||||
|
||||
@ -4923,77 +4921,6 @@ class Config:
|
||||
return self.parser.parse_args(self.all_args)
|
||||
|
||||
|
||||
class WebSocketsWrapper:
|
||||
"""Wraps websockets module to use in non-async scopes"""
|
||||
pool = None
|
||||
|
||||
def __init__(self, url, headers=None, connect=True):
|
||||
self.loop = asyncio.new_event_loop()
|
||||
# XXX: "loop" is deprecated
|
||||
self.conn = websockets.connect(
|
||||
url, extra_headers=headers, ping_interval=None,
|
||||
close_timeout=float('inf'), loop=self.loop, ping_timeout=float('inf'))
|
||||
if connect:
|
||||
self.__enter__()
|
||||
atexit.register(self.__exit__, None, None, None)
|
||||
|
||||
def __enter__(self):
|
||||
if not self.pool:
|
||||
self.pool = self.run_with_loop(self.conn.__aenter__(), self.loop)
|
||||
return self
|
||||
|
||||
def send(self, *args):
|
||||
self.run_with_loop(self.pool.send(*args), self.loop)
|
||||
|
||||
def recv(self, *args):
|
||||
return self.run_with_loop(self.pool.recv(*args), self.loop)
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
try:
|
||||
return self.run_with_loop(self.conn.__aexit__(type, value, traceback), self.loop)
|
||||
finally:
|
||||
self.loop.close()
|
||||
self._cancel_all_tasks(self.loop)
|
||||
|
||||
# taken from https://github.com/python/cpython/blob/3.9/Lib/asyncio/runners.py with modifications
|
||||
# for contributors: If there's any new library using asyncio needs to be run in non-async, move these function out of this class
|
||||
@staticmethod
|
||||
def run_with_loop(main, loop):
|
||||
if not asyncio.iscoroutine(main):
|
||||
raise ValueError(f'a coroutine was expected, got {main!r}')
|
||||
|
||||
try:
|
||||
return loop.run_until_complete(main)
|
||||
finally:
|
||||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||||
if hasattr(loop, 'shutdown_default_executor'):
|
||||
loop.run_until_complete(loop.shutdown_default_executor())
|
||||
|
||||
@staticmethod
|
||||
def _cancel_all_tasks(loop):
|
||||
to_cancel = asyncio.all_tasks(loop)
|
||||
|
||||
if not to_cancel:
|
||||
return
|
||||
|
||||
for task in to_cancel:
|
||||
task.cancel()
|
||||
|
||||
# XXX: "loop" is removed in python 3.10+
|
||||
loop.run_until_complete(
|
||||
asyncio.gather(*to_cancel, loop=loop, return_exceptions=True))
|
||||
|
||||
for task in to_cancel:
|
||||
if task.cancelled():
|
||||
continue
|
||||
if task.exception() is not None:
|
||||
loop.call_exception_handler({
|
||||
'message': 'unhandled exception during asyncio.run() shutdown',
|
||||
'exception': task.exception(),
|
||||
'task': task,
|
||||
})
|
||||
|
||||
|
||||
def merge_headers(*dicts):
|
||||
"""Merge dicts of http headers case insensitively, prioritizing the latter ones"""
|
||||
return {k.title(): v for k, v in itertools.chain.from_iterable(map(dict.items, dicts))}
|
||||
|
Reference in New Issue
Block a user