[utils] WebSockets wrapper for non-async functions (#2417)

Authored by: Lesmiscore

1 files changed, 69 insertions, 0 deletions

diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py
index bb8d65cad..c5489d494 100644
--- a/yt_dlp/utils.py
+++ b/yt_dlp/utils.py
@@ -3,6 +3,7 @@
 
 from __future__ import unicode_literals
 
+import asyncio
 import base64
 import binascii
 import calendar
@@ -73,6 +74,7 @@ from .compat import (
     compat_urllib_parse_unquote_plus,
     compat_urllib_request,
     compat_urlparse,
+    compat_websockets,
     compat_xpath,
 )
 
@@ -5311,3 +5313,70 @@ class Config:
 
     def parse_args(self):
         return self._parser.parse_args(list(self.all_args))
+
+
+class WebSocketsWrapper():
+    """Wraps websockets module to use in non-async scopes"""
+
+    def __init__(self, url, headers=None):
+        self.loop = asyncio.events.new_event_loop()
+        self.conn = compat_websockets.connect(
+            url, extra_headers=headers, ping_interval=None,
+            close_timeout=float('inf'), loop=self.loop, ping_timeout=float('inf'))
+
+    def __enter__(self):
+        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.r_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.coroutines.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.tasks.all_tasks(loop)
+
+        if not to_cancel:
+            return
+
+        for task in to_cancel:
+            task.cancel()
+
+        loop.run_until_complete(
+            asyncio.tasks.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,
+                })
+
+
+has_websockets = bool(compat_websockets)