From 3cea3edd1ac1101bd709dfa0305509028118b163 Mon Sep 17 00:00:00 2001
From: "Lesmiscore (Naoya Ozaki)" <nao20010128@gmail.com>
Date: Fri, 25 Mar 2022 17:24:39 +0900
Subject: [PATCH] [utils] WebSocketsWrapper: Allow omitting `__enter__`
invocation (#3187)
Authored by: Lesmiscore
---
yt_dlp/extractor/fc2.py | 1 -
yt_dlp/utils.py | 7 +++++--
2 files changed, 5 insertions(+), 3 deletions(-)
diff --git a/yt_dlp/extractor/fc2.py b/yt_dlp/extractor/fc2.py
index 7fc6b0e3d..54a83aa16 100644
--- a/yt_dlp/extractor/fc2.py
+++ b/yt_dlp/extractor/fc2.py
@@ -212,7 +212,6 @@ def _real_extract(self, url):
'Accept': '*/*',
'User-Agent': std_headers['User-Agent'],
})
- ws.__enter__()
self.write_debug('[debug] Sending HLS server request')
diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py
index 4de5f9626..e359c6bba 100644
--- a/yt_dlp/utils.py
+++ b/yt_dlp/utils.py
@@ -5436,15 +5436,18 @@ def parse_args(self):
class WebSocketsWrapper():
"""Wraps websockets module to use in non-async scopes"""
- def __init__(self, url, headers=None):
+ def __init__(self, url, headers=None, connect=True):
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'))
+ if connect:
+ self.__enter__()
atexit.register(self.__exit__, None, None, None)
def __enter__(self):
- self.pool = self.run_with_loop(self.conn.__aenter__(), self.loop)
+ if not self.pool:
+ self.pool = self.run_with_loop(self.conn.__aenter__(), self.loop)
return self
def send(self, *args):