From f0c9fb96827ff798a48626e7e5d32a9c5de7b97e Mon Sep 17 00:00:00 2001 From: pukkandan Date: Thu, 16 Jun 2022 02:25:43 +0530 Subject: [PATCH] [utils] `Popen`: Refactor to use contextmanager Fixes https://github.com/yt-dlp/yt-dlp/issues/3531#issuecomment-1156223597 --- yt_dlp/YoutubeDL.py | 12 ++--- yt_dlp/cookies.py | 33 +++++-------- yt_dlp/downloader/external.py | 66 +++++++++++--------------- yt_dlp/downloader/rtmp.py | 3 +- yt_dlp/extractor/openload.py | 13 ++--- yt_dlp/postprocessor/embedthumbnail.py | 10 ++-- yt_dlp/postprocessor/ffmpeg.py | 28 +++++------ yt_dlp/postprocessor/sponskrub.py | 14 +++--- yt_dlp/utils.py | 42 +++++++++------- 9 files changed, 98 insertions(+), 123 deletions(-) diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 1932af3fe..ffb0e1adf 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -3705,14 +3705,12 @@ def get_encoding(stream): if source == 'source': try: - sp = Popen( + stdout, _, _ = Popen.run( ['git', 'rev-parse', '--short', 'HEAD'], - stdout=subprocess.PIPE, stderr=subprocess.PIPE, - cwd=os.path.dirname(os.path.abspath(__file__))) - out, err = sp.communicate_or_kill() - out = out.decode().strip() - if re.match('[0-9a-f]+', out): - write_debug('Git HEAD: %s' % out) + text=True, cwd=os.path.dirname(os.path.abspath(__file__)), + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if re.fullmatch('[0-9a-f]+', stdout.strip()): + write_debug(f'Git HEAD: {stdout.strip()}') except Exception: with contextlib.suppress(Exception): sys.exc_clear() diff --git a/yt_dlp/cookies.py b/yt_dlp/cookies.py index 3978a6bf3..a74701750 100644 --- a/yt_dlp/cookies.py +++ b/yt_dlp/cookies.py @@ -709,21 +709,19 @@ def _get_kwallet_network_wallet(logger): """ default_wallet = 'kdewallet' try: - proc = Popen([ + stdout, _, returncode = Popen.run([ 'dbus-send', '--session', '--print-reply=literal', '--dest=org.kde.kwalletd5', '/modules/kwalletd5', '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 proc.returncode != 0: + if returncode: logger.warning('failed to read NetworkWallet') return default_wallet else: - network_wallet = stdout.decode().strip() - logger.debug(f'NetworkWallet = "{network_wallet}"') - return network_wallet + logger.debug(f'NetworkWallet = "{stdout.strip()}"') + return stdout.strip() except Exception as e: logger.warning(f'exception while obtaining NetworkWallet: {e}') return default_wallet @@ -741,17 +739,16 @@ def _get_kwallet_password(browser_keyring_name, logger): network_wallet = _get_kwallet_network_wallet(logger) try: - proc = Popen([ + stdout, _, returncode = Popen.run([ 'kwallet-query', '--read-password', f'{browser_keyring_name} Safe Storage', '--folder', f'{browser_keyring_name} Keys', network_wallet ], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) - stdout, stderr = proc.communicate_or_kill() - if proc.returncode != 0: - logger.error(f'kwallet-query failed with return code {proc.returncode}. Please consult ' - 'the kwallet-query man page for details') + if returncode: + logger.error(f'kwallet-query failed with return code {returncode}. ' + 'Please consult the kwallet-query man page for details') return b'' else: if stdout.lower().startswith(b'failed to read'): @@ -766,9 +763,7 @@ def _get_kwallet_password(browser_keyring_name, logger): return b'' else: logger.debug('password found') - if stdout[-1:] == b'\n': - stdout = stdout[:-1] - return stdout + return stdout.rstrip(b'\n') except Exception as e: logger.warning(f'exception running kwallet-query: {error_to_str(e)}') 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): logger.debug('using find-generic-password to obtain password from OSX keychain') try: - proc = Popen( + stdout, _, _ = Popen.run( ['security', 'find-generic-password', '-w', # write password to stdout '-a', browser_keyring_name, # match 'account' '-s', f'{browser_keyring_name} Safe Storage'], # match 'service' stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) - - stdout, stderr = proc.communicate_or_kill() - if stdout[-1:] == b'\n': - stdout = stdout[:-1] - return stdout + return stdout.rstrip(b'\n') except Exception as e: logger.warning(f'exception running find-generic-password: {error_to_str(e)}') return None diff --git a/yt_dlp/downloader/external.py b/yt_dlp/downloader/external.py index 3ef7fd4dc..a1cb07e05 100644 --- a/yt_dlp/downloader/external.py +++ b/yt_dlp/downloader/external.py @@ -34,6 +34,7 @@ class Features(enum.Enum): class ExternalFD(FragmentFD): SUPPORTED_PROTOCOLS = ('http', 'https', 'ftp', 'ftps') SUPPORTED_FEATURES = () + _CAPTURE_STDERR = True def real_download(self, filename, info_dict): self.report_destination(filename) @@ -128,24 +129,25 @@ def _call_downloader(self, tmpfilename, info_dict): self._debug_cmd(cmd) if 'fragments' not in info_dict: - p = Popen(cmd, stderr=subprocess.PIPE) - _, stderr = p.communicate_or_kill() - if p.returncode != 0: - self.to_stderr(stderr.decode('utf-8', 'replace')) - return p.returncode + _, stderr, returncode = Popen.run( + cmd, text=True, stderr=subprocess.PIPE if self._CAPTURE_STDERR else None) + if returncode and stderr: + self.to_stderr(stderr) + return returncode fragment_retries = self.params.get('fragment_retries', 0) skip_unavailable_fragments = self.params.get('skip_unavailable_fragments', True) count = 0 while count <= fragment_retries: - p = Popen(cmd, stderr=subprocess.PIPE) - _, stderr = p.communicate_or_kill() - if p.returncode == 0: + _, stderr, returncode = Popen.run(cmd, text=True, stderr=subprocess.PIPE) + if not returncode: break + # TODO: Decide whether to retry based on error code # 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 if count <= fragment_retries: self.to_screen( @@ -180,6 +182,7 @@ def _call_downloader(self, tmpfilename, info_dict): class CurlFD(ExternalFD): AVAILABLE_OPT = '-V' + _CAPTURE_STDERR = False # curl writes the progress to stderr def _make_cmd(self, tmpfilename, info_dict): cmd = [self.exe, '--location', '-o', tmpfilename, '--compressed'] @@ -204,16 +207,6 @@ def _make_cmd(self, tmpfilename, info_dict): cmd += ['--', info_dict['url']] 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): AVAILABLE_OPT = '-V' @@ -500,24 +493,23 @@ def _call_downloader(self, tmpfilename, info_dict): args.append(encodeFilename(ffpp._ffmpeg_filename_argument(tmpfilename), True)) self._debug_cmd(args) - proc = Popen(args, stdin=subprocess.PIPE, env=env) - if url in ('-', 'pipe:'): - self.on_process_started(proc, proc.stdin) - try: - retval = proc.wait() - except BaseException as e: - # 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 - # produces a file that is playable (this is mostly useful for live - # streams). Note that Windows is not affected and produces playable - # files (see https://github.com/ytdl-org/youtube-dl/issues/8300). - if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32' and url not in ('-', 'pipe:'): - proc.communicate_or_kill(b'q') - else: - proc.kill() - proc.wait() - raise - return retval + with Popen(args, stdin=subprocess.PIPE, env=env) as proc: + if url in ('-', 'pipe:'): + self.on_process_started(proc, proc.stdin) + try: + retval = proc.wait() + except BaseException as e: + # 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 + # produces a file that is playable (this is mostly useful for live + # streams). Note that Windows is not affected and produces playable + # files (see https://github.com/ytdl-org/youtube-dl/issues/8300). + if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32' and url not in ('-', 'pipe:'): + proc.communicate_or_kill(b'q') + else: + proc.kill(timeout=None) + raise + return retval class AVconvFD(FFmpegFD): diff --git a/yt_dlp/downloader/rtmp.py b/yt_dlp/downloader/rtmp.py index 3464eeef9..217158952 100644 --- a/yt_dlp/downloader/rtmp.py +++ b/yt_dlp/downloader/rtmp.py @@ -92,8 +92,7 @@ def run_rtmpdump(args): self.to_screen('') return proc.wait() except BaseException: # Including KeyboardInterrupt - proc.kill() - proc.wait() + proc.kill(timeout=None) raise url = info_dict['url'] diff --git a/yt_dlp/extractor/openload.py b/yt_dlp/extractor/openload.py index 61e3a8b86..d987cd927 100644 --- a/yt_dlp/extractor/openload.py +++ b/yt_dlp/extractor/openload.py @@ -9,7 +9,6 @@ ExtractorError, Popen, check_executable, - encodeArgument, get_exe_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: 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], - stdout=subprocess.PIPE, stderr=subprocess.PIPE) - out, err = p.communicate_or_kill() - if p.returncode != 0: - raise ExtractorError( - 'Executing JS failed\n:' + encodeArgument(err)) + text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if returncode: + raise ExtractorError(f'Executing JS failed\n:{stderr}') with open(self._TMP_FILES['html'].name, 'rb') as f: html = f.read().decode('utf-8') self._load_cookies() - return (html, encodeArgument(out)) + return (html, stdout) diff --git a/yt_dlp/postprocessor/embedthumbnail.py b/yt_dlp/postprocessor/embedthumbnail.py index e031d344f..606d90d3d 100644 --- a/yt_dlp/postprocessor/embedthumbnail.py +++ b/yt_dlp/postprocessor/embedthumbnail.py @@ -157,14 +157,12 @@ def run(self, info): self._report_run('atomicparsley', filename) self.write_debug('AtomicParsley command line: %s' % shell_quote(cmd)) - p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - stdout, stderr = p.communicate_or_kill() - if p.returncode != 0: - msg = stderr.decode('utf-8', 'replace').strip() - self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {msg}') + stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if returncode: + self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {stderr.strip()}') # for formats that don't support thumbnails (like 3gp) AtomicParsley # 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') success = False diff --git a/yt_dlp/postprocessor/ffmpeg.py b/yt_dlp/postprocessor/ffmpeg.py index a726a962f..71ae16b51 100644 --- a/yt_dlp/postprocessor/ffmpeg.py +++ b/yt_dlp/postprocessor/ffmpeg.py @@ -239,14 +239,12 @@ def get_audio_codec(self, path): encodeArgument('-i')] cmd.append(encodeFilename(self._ffmpeg_filename_argument(path), True)) self.write_debug(f'{self.basename} command line: {shell_quote(cmd)}') - handle = Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - stdout_data, stderr_data = handle.communicate_or_kill() - expected_ret = 0 if self.probe_available else 1 - if handle.wait() != expected_ret: + stdout, stderr, returncode = Popen.run(cmd, text=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if returncode != (0 if self.probe_available else 1): return None except OSError: 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: audio_codec = None for line in output.split('\n'): @@ -280,11 +278,10 @@ def get_metadata_object(self, path, opts=[]): ] cmd += opts - cmd.append(encodeFilename(self._ffmpeg_filename_argument(path), True)) - self.write_debug('ffprobe command line: %s' % shell_quote(cmd)) - p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) - stdout, stderr = p.communicate() - return json.loads(stdout.decode('utf-8', 'replace')) + cmd.append(self._ffmpeg_filename_argument(path)) + self.write_debug(f'ffprobe command line: {shell_quote(cmd)}') + stdout, _, _ = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) + return json.loads(stdout) def get_stream_number(self, path, keys, value): 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) self.write_debug('ffmpeg command line: %s' % shell_quote(cmd)) - p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) - stdout, stderr = p.communicate_or_kill() - if p.returncode not in variadic(expected_retcodes): - stderr = stderr.decode('utf-8', 'replace').strip() - self.write_debug(stderr) - raise FFmpegPostProcessorError(stderr.split('\n')[-1]) + stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) + if returncode not in variadic(expected_retcodes): + raise FFmpegPostProcessorError(stderr.strip().splitlines()[-1]) for out_path, _ in output_path_opts: if out_path: 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): return self.run_ffmpeg_multiple_files([path], out_path, opts, **kwargs) diff --git a/yt_dlp/postprocessor/sponskrub.py b/yt_dlp/postprocessor/sponskrub.py index 1a9f5dc66..ff50d5b4f 100644 --- a/yt_dlp/postprocessor/sponskrub.py +++ b/yt_dlp/postprocessor/sponskrub.py @@ -84,17 +84,15 @@ def run(self, information): cmd = [encodeArgument(i) for i in cmd] self.write_debug('sponskrub command line: %s' % shell_quote(cmd)) - pipe = None if self.get_param('verbose') else subprocess.PIPE - p = Popen(cmd, stdout=pipe) - stdout = p.communicate_or_kill()[0] + stdout, _, returncode = Popen.run(cmd, text=True, stdout=None if self.get_param('verbose') else subprocess.PIPE) - if p.returncode == 0: + if not returncode: os.replace(temp_filename, filename) 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') else: - msg = stdout.decode('utf-8', 'replace').strip() if stdout else '' - msg = msg.split('\n')[0 if msg.lower().startswith('unrecognised') else -1] - raise PostProcessingError(msg if msg else 'sponskrub failed with error code %s' % p.returncode) + raise PostProcessingError( + stdout.strip().splitlines()[0 if stdout.strip().lower().startswith('unrecognised') else -1] + or f'sponskrub failed with error code {returncode}') return [], information diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py index 11ef7744c..be7cbf9fd 100644 --- a/yt_dlp/utils.py +++ b/yt_dlp/utils.py @@ -841,17 +841,31 @@ class Popen(subprocess.Popen): else: _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) def communicate_or_kill(self, *args, **kwargs): try: return self.communicate(*args, **kwargs) except BaseException: # Including KeyboardInterrupt - self.kill() - self.wait() + self.kill(timeout=None) 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(): 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. args can be a list of arguments for a short output (like -version) """ 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: return False 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 # SIGTTOU if yt-dlp is run in the background. # See https://github.com/ytdl-org/youtube-dl/issues/955#issuecomment-209789656 - out, _ = Popen( - [encodeArgument(exe)] + args, stdin=subprocess.PIPE, - stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate_or_kill() + stdout, _, _ = Popen.run([encodeArgument(exe)] + args, text=True, + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) except OSError: return False - if isinstance(out, bytes): # Python 2.x - out = out.decode('ascii', 'ignore') - return out + return stdout def detect_exe_version(output, version_re=None, unrecognized='present'): @@ -4796,14 +4807,13 @@ def write_xattr(path, key, value): value = value.decode() try: - p = Popen( + _, stderr, returncode = Popen.run( [exe, '-w', key, value, path] if exe == 'xattr' else [exe, '-n', key, '-v', value, path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) except OSError as e: raise XAttrMetadataError(e.errno, e.strerror) - stderr = p.communicate_or_kill()[1].decode('utf-8', 'replace') - if p.returncode: - raise XAttrMetadataError(p.returncode, stderr) + if returncode: + raise XAttrMetadataError(returncode, stderr) 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): return global WINDOWS_VT_MODE - startupinfo = subprocess.STARTUPINFO() - startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW try: - subprocess.Popen('', shell=True, startupinfo=startupinfo).wait() + Popen.run('', shell=True) except Exception: return