1
0
mirror of https://github.com/yt-dlp/yt-dlp.git synced 2024-11-02 17:22:31 +01:00

[utils] WebSocketsWrapper: Allow omitting __enter__ invocation (#3187)

Authored by: Lesmiscore
This commit is contained in:
Lesmiscore (Naoya Ozaki) 2022-03-25 17:24:39 +09:00 committed by GitHub
parent b1a7cd056a
commit 3cea3edd1a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 5 additions and 3 deletions

View File

@ -212,7 +212,6 @@ def _real_extract(self, url):
'Accept': '*/*', 'Accept': '*/*',
'User-Agent': std_headers['User-Agent'], 'User-Agent': std_headers['User-Agent'],
}) })
ws.__enter__()
self.write_debug('[debug] Sending HLS server request') self.write_debug('[debug] Sending HLS server request')

View File

@ -5436,14 +5436,17 @@ def parse_args(self):
class WebSocketsWrapper(): class WebSocketsWrapper():
"""Wraps websockets module to use in non-async scopes""" """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.loop = asyncio.events.new_event_loop()
self.conn = compat_websockets.connect( self.conn = compat_websockets.connect(
url, extra_headers=headers, ping_interval=None, url, extra_headers=headers, ping_interval=None,
close_timeout=float('inf'), loop=self.loop, ping_timeout=float('inf')) close_timeout=float('inf'), loop=self.loop, ping_timeout=float('inf'))
if connect:
self.__enter__()
atexit.register(self.__exit__, None, None, None) atexit.register(self.__exit__, None, None, None)
def __enter__(self): def __enter__(self):
if not self.pool:
self.pool = self.run_with_loop(self.conn.__aenter__(), self.loop) self.pool = self.run_with_loop(self.conn.__aenter__(), self.loop)
return self return self