[utils] Popen: Refactor to use contextmanager

Fixes https://github.com/yt-dlp/yt-dlp/issues/3531#issuecomment-1156223597
This commit is contained in:
pukkandan 2022-06-16 02:25:43 +05:30
parent 560738f34d
commit f0c9fb9682
No known key found for this signature in database
GPG key ID: 7EEE9E1E817D0A39
9 changed files with 98 additions and 123 deletions

View file

@ -3705,14 +3705,12 @@ def get_encoding(stream):
if source == 'source': if source == 'source':
try: try:
sp = Popen( stdout, _, _ = Popen.run(
['git', 'rev-parse', '--short', 'HEAD'], ['git', 'rev-parse', '--short', 'HEAD'],
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, cwd=os.path.dirname(os.path.abspath(__file__)),
cwd=os.path.dirname(os.path.abspath(__file__))) stdout=subprocess.PIPE, stderr=subprocess.PIPE)
out, err = sp.communicate_or_kill() if re.fullmatch('[0-9a-f]+', stdout.strip()):
out = out.decode().strip() write_debug(f'Git HEAD: {stdout.strip()}')
if re.match('[0-9a-f]+', out):
write_debug('Git HEAD: %s' % out)
except Exception: except Exception:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
sys.exc_clear() sys.exc_clear()

View file

@ -709,21 +709,19 @@ def _get_kwallet_network_wallet(logger):
""" """
default_wallet = 'kdewallet' default_wallet = 'kdewallet'
try: try:
proc = Popen([ stdout, _, returncode = Popen.run([
'dbus-send', '--session', '--print-reply=literal', 'dbus-send', '--session', '--print-reply=literal',
'--dest=org.kde.kwalletd5', '--dest=org.kde.kwalletd5',
'/modules/kwalletd5', '/modules/kwalletd5',
'org.kde.KWallet.networkWallet' 'org.kde.KWallet.networkWallet'
], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) ], text=True, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
stdout, stderr = proc.communicate_or_kill() if returncode:
if proc.returncode != 0:
logger.warning('failed to read NetworkWallet') logger.warning('failed to read NetworkWallet')
return default_wallet return default_wallet
else: else:
network_wallet = stdout.decode().strip() logger.debug(f'NetworkWallet = "{stdout.strip()}"')
logger.debug(f'NetworkWallet = "{network_wallet}"') return stdout.strip()
return network_wallet
except Exception as e: except Exception as e:
logger.warning(f'exception while obtaining NetworkWallet: {e}') logger.warning(f'exception while obtaining NetworkWallet: {e}')
return default_wallet return default_wallet
@ -741,17 +739,16 @@ def _get_kwallet_password(browser_keyring_name, logger):
network_wallet = _get_kwallet_network_wallet(logger) network_wallet = _get_kwallet_network_wallet(logger)
try: try:
proc = Popen([ stdout, _, returncode = Popen.run([
'kwallet-query', 'kwallet-query',
'--read-password', f'{browser_keyring_name} Safe Storage', '--read-password', f'{browser_keyring_name} Safe Storage',
'--folder', f'{browser_keyring_name} Keys', '--folder', f'{browser_keyring_name} Keys',
network_wallet network_wallet
], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) ], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
stdout, stderr = proc.communicate_or_kill() if returncode:
if proc.returncode != 0: logger.error(f'kwallet-query failed with return code {returncode}. '
logger.error(f'kwallet-query failed with return code {proc.returncode}. Please consult ' 'Please consult the kwallet-query man page for details')
'the kwallet-query man page for details')
return b'' return b''
else: else:
if stdout.lower().startswith(b'failed to read'): if stdout.lower().startswith(b'failed to read'):
@ -766,9 +763,7 @@ def _get_kwallet_password(browser_keyring_name, logger):
return b'' return b''
else: else:
logger.debug('password found') logger.debug('password found')
if stdout[-1:] == b'\n': return stdout.rstrip(b'\n')
stdout = stdout[:-1]
return stdout
except Exception as e: except Exception as e:
logger.warning(f'exception running kwallet-query: {error_to_str(e)}') logger.warning(f'exception running kwallet-query: {error_to_str(e)}')
return b'' return b''
@ -815,17 +810,13 @@ def _get_linux_keyring_password(browser_keyring_name, keyring, logger):
def _get_mac_keyring_password(browser_keyring_name, logger): def _get_mac_keyring_password(browser_keyring_name, logger):
logger.debug('using find-generic-password to obtain password from OSX keychain') logger.debug('using find-generic-password to obtain password from OSX keychain')
try: try:
proc = Popen( stdout, _, _ = Popen.run(
['security', 'find-generic-password', ['security', 'find-generic-password',
'-w', # write password to stdout '-w', # write password to stdout
'-a', browser_keyring_name, # match 'account' '-a', browser_keyring_name, # match 'account'
'-s', f'{browser_keyring_name} Safe Storage'], # match 'service' '-s', f'{browser_keyring_name} Safe Storage'], # match 'service'
stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
return stdout.rstrip(b'\n')
stdout, stderr = proc.communicate_or_kill()
if stdout[-1:] == b'\n':
stdout = stdout[:-1]
return stdout
except Exception as e: except Exception as e:
logger.warning(f'exception running find-generic-password: {error_to_str(e)}') logger.warning(f'exception running find-generic-password: {error_to_str(e)}')
return None return None

View file

@ -34,6 +34,7 @@ class Features(enum.Enum):
class ExternalFD(FragmentFD): class ExternalFD(FragmentFD):
SUPPORTED_PROTOCOLS = ('http', 'https', 'ftp', 'ftps') SUPPORTED_PROTOCOLS = ('http', 'https', 'ftp', 'ftps')
SUPPORTED_FEATURES = () SUPPORTED_FEATURES = ()
_CAPTURE_STDERR = True
def real_download(self, filename, info_dict): def real_download(self, filename, info_dict):
self.report_destination(filename) self.report_destination(filename)
@ -128,24 +129,25 @@ def _call_downloader(self, tmpfilename, info_dict):
self._debug_cmd(cmd) self._debug_cmd(cmd)
if 'fragments' not in info_dict: if 'fragments' not in info_dict:
p = Popen(cmd, stderr=subprocess.PIPE) _, stderr, returncode = Popen.run(
_, stderr = p.communicate_or_kill() cmd, text=True, stderr=subprocess.PIPE if self._CAPTURE_STDERR else None)
if p.returncode != 0: if returncode and stderr:
self.to_stderr(stderr.decode('utf-8', 'replace')) self.to_stderr(stderr)
return p.returncode return returncode
fragment_retries = self.params.get('fragment_retries', 0) fragment_retries = self.params.get('fragment_retries', 0)
skip_unavailable_fragments = self.params.get('skip_unavailable_fragments', True) skip_unavailable_fragments = self.params.get('skip_unavailable_fragments', True)
count = 0 count = 0
while count <= fragment_retries: while count <= fragment_retries:
p = Popen(cmd, stderr=subprocess.PIPE) _, stderr, returncode = Popen.run(cmd, text=True, stderr=subprocess.PIPE)
_, stderr = p.communicate_or_kill() if not returncode:
if p.returncode == 0:
break break
# TODO: Decide whether to retry based on error code # TODO: Decide whether to retry based on error code
# https://aria2.github.io/manual/en/html/aria2c.html#exit-status # https://aria2.github.io/manual/en/html/aria2c.html#exit-status
self.to_stderr(stderr.decode('utf-8', 'replace')) if stderr:
self.to_stderr(stderr)
count += 1 count += 1
if count <= fragment_retries: if count <= fragment_retries:
self.to_screen( self.to_screen(
@ -180,6 +182,7 @@ def _call_downloader(self, tmpfilename, info_dict):
class CurlFD(ExternalFD): class CurlFD(ExternalFD):
AVAILABLE_OPT = '-V' AVAILABLE_OPT = '-V'
_CAPTURE_STDERR = False # curl writes the progress to stderr
def _make_cmd(self, tmpfilename, info_dict): def _make_cmd(self, tmpfilename, info_dict):
cmd = [self.exe, '--location', '-o', tmpfilename, '--compressed'] cmd = [self.exe, '--location', '-o', tmpfilename, '--compressed']
@ -204,16 +207,6 @@ def _make_cmd(self, tmpfilename, info_dict):
cmd += ['--', info_dict['url']] cmd += ['--', info_dict['url']]
return cmd return cmd
def _call_downloader(self, tmpfilename, info_dict):
cmd = [encodeArgument(a) for a in self._make_cmd(tmpfilename, info_dict)]
self._debug_cmd(cmd)
# curl writes the progress to stderr so don't capture it.
p = Popen(cmd)
p.communicate_or_kill()
return p.returncode
class AxelFD(ExternalFD): class AxelFD(ExternalFD):
AVAILABLE_OPT = '-V' AVAILABLE_OPT = '-V'
@ -500,24 +493,23 @@ def _call_downloader(self, tmpfilename, info_dict):
args.append(encodeFilename(ffpp._ffmpeg_filename_argument(tmpfilename), True)) args.append(encodeFilename(ffpp._ffmpeg_filename_argument(tmpfilename), True))
self._debug_cmd(args) self._debug_cmd(args)
proc = Popen(args, stdin=subprocess.PIPE, env=env) with Popen(args, stdin=subprocess.PIPE, env=env) as proc:
if url in ('-', 'pipe:'): if url in ('-', 'pipe:'):
self.on_process_started(proc, proc.stdin) self.on_process_started(proc, proc.stdin)
try: try:
retval = proc.wait() retval = proc.wait()
except BaseException as e: except BaseException as e:
# subprocces.run would send the SIGKILL signal to ffmpeg and the # subprocces.run would send the SIGKILL signal to ffmpeg and the
# mp4 file couldn't be played, but if we ask ffmpeg to quit it # mp4 file couldn't be played, but if we ask ffmpeg to quit it
# produces a file that is playable (this is mostly useful for live # produces a file that is playable (this is mostly useful for live
# streams). Note that Windows is not affected and produces playable # streams). Note that Windows is not affected and produces playable
# files (see https://github.com/ytdl-org/youtube-dl/issues/8300). # files (see https://github.com/ytdl-org/youtube-dl/issues/8300).
if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32' and url not in ('-', 'pipe:'): if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32' and url not in ('-', 'pipe:'):
proc.communicate_or_kill(b'q') proc.communicate_or_kill(b'q')
else: else:
proc.kill() proc.kill(timeout=None)
proc.wait() raise
raise return retval
return retval
class AVconvFD(FFmpegFD): class AVconvFD(FFmpegFD):

View file

@ -92,8 +92,7 @@ def run_rtmpdump(args):
self.to_screen('') self.to_screen('')
return proc.wait() return proc.wait()
except BaseException: # Including KeyboardInterrupt except BaseException: # Including KeyboardInterrupt
proc.kill() proc.kill(timeout=None)
proc.wait()
raise raise
url = info_dict['url'] url = info_dict['url']

View file

@ -9,7 +9,6 @@
ExtractorError, ExtractorError,
Popen, Popen,
check_executable, check_executable,
encodeArgument,
get_exe_version, get_exe_version,
is_outdated_version, is_outdated_version,
) )
@ -213,16 +212,14 @@ def get(self, url, html=None, video_id=None, note=None, note2='Executing JS on w
else: else:
self.extractor.to_screen(f'{video_id}: {note2}') self.extractor.to_screen(f'{video_id}: {note2}')
p = Popen( stdout, stderr, returncode = Popen.run(
[self.exe, '--ssl-protocol=any', self._TMP_FILES['script'].name], [self.exe, '--ssl-protocol=any', self._TMP_FILES['script'].name],
stdout=subprocess.PIPE, stderr=subprocess.PIPE) text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
out, err = p.communicate_or_kill() if returncode:
if p.returncode != 0: raise ExtractorError(f'Executing JS failed\n:{stderr}')
raise ExtractorError(
'Executing JS failed\n:' + encodeArgument(err))
with open(self._TMP_FILES['html'].name, 'rb') as f: with open(self._TMP_FILES['html'].name, 'rb') as f:
html = f.read().decode('utf-8') html = f.read().decode('utf-8')
self._load_cookies() self._load_cookies()
return (html, encodeArgument(out)) return (html, stdout)

View file

@ -157,14 +157,12 @@ def run(self, info):
self._report_run('atomicparsley', filename) self._report_run('atomicparsley', filename)
self.write_debug('AtomicParsley command line: %s' % shell_quote(cmd)) self.write_debug('AtomicParsley command line: %s' % shell_quote(cmd))
p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = p.communicate_or_kill() if returncode:
if p.returncode != 0: self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {stderr.strip()}')
msg = stderr.decode('utf-8', 'replace').strip()
self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {msg}')
# for formats that don't support thumbnails (like 3gp) AtomicParsley # for formats that don't support thumbnails (like 3gp) AtomicParsley
# won't create to the temporary file # won't create to the temporary file
if b'No changes' in stdout: if 'No changes' in stdout:
self.report_warning('The file format doesn\'t support embedding a thumbnail') self.report_warning('The file format doesn\'t support embedding a thumbnail')
success = False success = False

View file

@ -239,14 +239,12 @@ def get_audio_codec(self, path):
encodeArgument('-i')] encodeArgument('-i')]
cmd.append(encodeFilename(self._ffmpeg_filename_argument(path), True)) cmd.append(encodeFilename(self._ffmpeg_filename_argument(path), True))
self.write_debug(f'{self.basename} command line: {shell_quote(cmd)}') self.write_debug(f'{self.basename} command line: {shell_quote(cmd)}')
handle = Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr, returncode = Popen.run(cmd, text=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout_data, stderr_data = handle.communicate_or_kill() if returncode != (0 if self.probe_available else 1):
expected_ret = 0 if self.probe_available else 1
if handle.wait() != expected_ret:
return None return None
except OSError: except OSError:
return None return None
output = (stdout_data if self.probe_available else stderr_data).decode('ascii', 'ignore') output = stdout if self.probe_available else stderr
if self.probe_available: if self.probe_available:
audio_codec = None audio_codec = None
for line in output.split('\n'): for line in output.split('\n'):
@ -280,11 +278,10 @@ def get_metadata_object(self, path, opts=[]):
] ]
cmd += opts cmd += opts
cmd.append(encodeFilename(self._ffmpeg_filename_argument(path), True)) cmd.append(self._ffmpeg_filename_argument(path))
self.write_debug('ffprobe command line: %s' % shell_quote(cmd)) self.write_debug(f'ffprobe command line: {shell_quote(cmd)}')
p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) stdout, _, _ = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
stdout, stderr = p.communicate() return json.loads(stdout)
return json.loads(stdout.decode('utf-8', 'replace'))
def get_stream_number(self, path, keys, value): def get_stream_number(self, path, keys, value):
streams = self.get_metadata_object(path)['streams'] streams = self.get_metadata_object(path)['streams']
@ -346,16 +343,13 @@ def make_args(file, args, name, number):
for i, (path, opts) in enumerate(path_opts) if path) for i, (path, opts) in enumerate(path_opts) if path)
self.write_debug('ffmpeg command line: %s' % shell_quote(cmd)) self.write_debug('ffmpeg command line: %s' % shell_quote(cmd))
p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
stdout, stderr = p.communicate_or_kill() if returncode not in variadic(expected_retcodes):
if p.returncode not in variadic(expected_retcodes): raise FFmpegPostProcessorError(stderr.strip().splitlines()[-1])
stderr = stderr.decode('utf-8', 'replace').strip()
self.write_debug(stderr)
raise FFmpegPostProcessorError(stderr.split('\n')[-1])
for out_path, _ in output_path_opts: for out_path, _ in output_path_opts:
if out_path: if out_path:
self.try_utime(out_path, oldest_mtime, oldest_mtime) self.try_utime(out_path, oldest_mtime, oldest_mtime)
return stderr.decode('utf-8', 'replace') return stderr
def run_ffmpeg(self, path, out_path, opts, **kwargs): def run_ffmpeg(self, path, out_path, opts, **kwargs):
return self.run_ffmpeg_multiple_files([path], out_path, opts, **kwargs) return self.run_ffmpeg_multiple_files([path], out_path, opts, **kwargs)

View file

@ -84,17 +84,15 @@ def run(self, information):
cmd = [encodeArgument(i) for i in cmd] cmd = [encodeArgument(i) for i in cmd]
self.write_debug('sponskrub command line: %s' % shell_quote(cmd)) self.write_debug('sponskrub command line: %s' % shell_quote(cmd))
pipe = None if self.get_param('verbose') else subprocess.PIPE stdout, _, returncode = Popen.run(cmd, text=True, stdout=None if self.get_param('verbose') else subprocess.PIPE)
p = Popen(cmd, stdout=pipe)
stdout = p.communicate_or_kill()[0]
if p.returncode == 0: if not returncode:
os.replace(temp_filename, filename) os.replace(temp_filename, filename)
self.to_screen('Sponsor sections have been %s' % ('removed' if self.cutout else 'marked')) self.to_screen('Sponsor sections have been %s' % ('removed' if self.cutout else 'marked'))
elif p.returncode == 3: elif returncode == 3:
self.to_screen('No segments in the SponsorBlock database') self.to_screen('No segments in the SponsorBlock database')
else: else:
msg = stdout.decode('utf-8', 'replace').strip() if stdout else '' raise PostProcessingError(
msg = msg.split('\n')[0 if msg.lower().startswith('unrecognised') else -1] stdout.strip().splitlines()[0 if stdout.strip().lower().startswith('unrecognised') else -1]
raise PostProcessingError(msg if msg else 'sponskrub failed with error code %s' % p.returncode) or f'sponskrub failed with error code {returncode}')
return [], information return [], information

View file

@ -841,17 +841,31 @@ class Popen(subprocess.Popen):
else: else:
_startupinfo = None _startupinfo = None
def __init__(self, *args, **kwargs): def __init__(self, *args, text=False, **kwargs):
if text is True:
kwargs['universal_newlines'] = True # For 3.6 compatibility
kwargs.setdefault('encoding', 'utf-8')
kwargs.setdefault('errors', 'replace')
super().__init__(*args, **kwargs, startupinfo=self._startupinfo) super().__init__(*args, **kwargs, startupinfo=self._startupinfo)
def communicate_or_kill(self, *args, **kwargs): def communicate_or_kill(self, *args, **kwargs):
try: try:
return self.communicate(*args, **kwargs) return self.communicate(*args, **kwargs)
except BaseException: # Including KeyboardInterrupt except BaseException: # Including KeyboardInterrupt
self.kill() self.kill(timeout=None)
self.wait()
raise raise
def kill(self, *, timeout=0):
super().kill()
if timeout != 0:
self.wait(timeout=timeout)
@classmethod
def run(cls, *args, **kwargs):
with cls(*args, **kwargs) as proc:
stdout, stderr = proc.communicate_or_kill()
return stdout or '', stderr or '', proc.returncode
def get_subprocess_encoding(): def get_subprocess_encoding():
if sys.platform == 'win32' and sys.getwindowsversion()[0] >= 5: if sys.platform == 'win32' and sys.getwindowsversion()[0] >= 5:
@ -2556,7 +2570,7 @@ def check_executable(exe, args=[]):
""" Checks if the given binary is installed somewhere in PATH, and returns its name. """ Checks if the given binary is installed somewhere in PATH, and returns its name.
args can be a list of arguments for a short output (like -version) """ args can be a list of arguments for a short output (like -version) """
try: try:
Popen([exe] + args, stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate_or_kill() Popen.run([exe] + args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
except OSError: except OSError:
return False return False
return exe return exe
@ -2569,14 +2583,11 @@ def _get_exe_version_output(exe, args, *, to_screen=None):
# STDIN should be redirected too. On UNIX-like systems, ffmpeg triggers # STDIN should be redirected too. On UNIX-like systems, ffmpeg triggers
# SIGTTOU if yt-dlp is run in the background. # SIGTTOU if yt-dlp is run in the background.
# See https://github.com/ytdl-org/youtube-dl/issues/955#issuecomment-209789656 # See https://github.com/ytdl-org/youtube-dl/issues/955#issuecomment-209789656
out, _ = Popen( stdout, _, _ = Popen.run([encodeArgument(exe)] + args, text=True,
[encodeArgument(exe)] + args, stdin=subprocess.PIPE, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate_or_kill()
except OSError: except OSError:
return False return False
if isinstance(out, bytes): # Python 2.x return stdout
out = out.decode('ascii', 'ignore')
return out
def detect_exe_version(output, version_re=None, unrecognized='present'): def detect_exe_version(output, version_re=None, unrecognized='present'):
@ -4796,14 +4807,13 @@ def write_xattr(path, key, value):
value = value.decode() value = value.decode()
try: try:
p = Popen( _, stderr, returncode = Popen.run(
[exe, '-w', key, value, path] if exe == 'xattr' else [exe, '-n', key, '-v', value, path], [exe, '-w', key, value, path] if exe == 'xattr' else [exe, '-n', key, '-v', value, path],
stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
except OSError as e: except OSError as e:
raise XAttrMetadataError(e.errno, e.strerror) raise XAttrMetadataError(e.errno, e.strerror)
stderr = p.communicate_or_kill()[1].decode('utf-8', 'replace') if returncode:
if p.returncode: raise XAttrMetadataError(returncode, stderr)
raise XAttrMetadataError(p.returncode, stderr)
def random_birthday(year_field, month_field, day_field): def random_birthday(year_field, month_field, day_field):
@ -5146,10 +5156,8 @@ def windows_enable_vt_mode(): # TODO: Do this the proper way https://bugs.pytho
if get_windows_version() < (10, 0, 10586): if get_windows_version() < (10, 0, 10586):
return return
global WINDOWS_VT_MODE global WINDOWS_VT_MODE
startupinfo = subprocess.STARTUPINFO()
startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
try: try:
subprocess.Popen('', shell=True, startupinfo=startupinfo).wait() Popen.run('', shell=True)
except Exception: except Exception:
return return