From e63878e23588f630bd11e7595b6ee738decbe4ed Mon Sep 17 00:00:00 2001 From: Simon Sawicki Date: Fri, 27 Dec 2024 16:17:45 +0100 Subject: [PATCH] Create a dedicated progress state class --- yt_dlp/YoutubeDL.py | 24 +++++++++++------------- yt_dlp/downloader/common.py | 13 ++----------- yt_dlp/postprocessor/common.py | 13 ++----------- yt_dlp/utils/_utils.py | 30 ++++++++++++++++++++++++++++++ 4 files changed, 45 insertions(+), 35 deletions(-) diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 4eb851592..7dae52293 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -157,7 +157,7 @@ write_json_file, write_string, ) -from .utils._utils import _UnsafeExtensionError, _YDLLogger +from .utils._utils import _UnsafeExtensionError, _YDLLogger, _ProgressState from .utils.networking import ( HTTPHeaderDict, clean_headers, @@ -955,14 +955,18 @@ def _send_console_code(self, code): self._write_string(code, self._out_files.console) return True - def to_console_title(self, message): + def to_console_title(self, message=None, progress_state=None, percent=None): if not self.params.get('consoletitle'): return - message = remove_terminal_sequences(message) - if not self._send_console_code(f'\033]0;{message}\007'): - if os.name == 'nt' and ctypes.windll.kernel32.GetConsoleWindow(): + + if message: + success = self._send_console_code(f'\033]0;{remove_terminal_sequences(message)}\007') + if not success and os.name == 'nt' and ctypes.windll.kernel32.GetConsoleWindow(): ctypes.windll.kernel32.SetConsoleTitleW(message) + if isinstance(progress_state, _ProgressState): + self._send_console_code(progress_state.get_ansi_escape(percent)) + def save_console_title(self): if not self.params.get('consoletitle') or self.params.get('simulate'): return @@ -975,10 +979,7 @@ def restore_console_title(self): def __enter__(self): self.save_console_title() - if self.params.get('consoletitle'): - # Set progress to "indeterminate" - # See: https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC - self._send_console_code('\033]9;4;3;0\007') + self.to_console_title(progress_state=_ProgressState.INDETERMINATE) return self def save_cookies(self): @@ -987,10 +988,7 @@ def save_cookies(self): def __exit__(self, *args): self.restore_console_title() - if self.params.get('consoletitle'): - # Set progress to "disabled" - # See: https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC - self._send_console_code('\033]9;4;0;0\007') + self.to_console_title(progress_state=_ProgressState.HIDDEN) self.close() def close(self): diff --git a/yt_dlp/downloader/common.py b/yt_dlp/downloader/common.py index 221f3e860..3ac8a7c63 100644 --- a/yt_dlp/downloader/common.py +++ b/yt_dlp/downloader/common.py @@ -31,6 +31,7 @@ timetuple_from_msec, try_call, ) +from ..utils._utils import _ProgressState class FileDownloader: @@ -333,17 +334,7 @@ def _report_progress_status(self, s, default_template): progress_dict), s.get('progress_idx') or 0) self.to_console_title(self.ydl.evaluate_outtmpl( progress_template.get('download-title') or 'yt-dlp %(progress._default_template)s', - progress_dict)) - - percent = s.get('_percent') - if s['status'] not in ('downloading', 'error', 'finished') or percent is None: - return - # Emit ConEmu progress codes: https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC - if s['status'] == 'finished': - self.ydl._send_console_code('\033]9;4;3;0\007') - return - state = 1 if s['status'] == 'downloading' else 2 - self.ydl._send_console_code(f'\033]9;4;{state};{int(percent)}\007') + progress_dict), _ProgressState.from_dict(s), s.get('_progress')) def _format_progress(self, *args, **kwargs): return self.ydl._format_text( diff --git a/yt_dlp/postprocessor/common.py b/yt_dlp/postprocessor/common.py index da40e0ff6..b547c6305 100644 --- a/yt_dlp/postprocessor/common.py +++ b/yt_dlp/postprocessor/common.py @@ -10,6 +10,7 @@ _configuration_args, deprecation_warning, ) +from ..utils._utils import _ProgressState class PostProcessorMetaClass(type): @@ -189,17 +190,7 @@ def report_progress(self, s): self._downloader.to_console_title(self._downloader.evaluate_outtmpl( progress_template.get('postprocess-title') or 'yt-dlp %(progress._default_template)s', - progress_dict)) - - percent = s.get('_percent') - if s['status'] not in ('downloading', 'error', 'finished') or percent is None: - return - # Emit ConEmu progress codes: https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC - if s['status'] == 'finished': - self._downloader._send_console_code('\033]9;4;3;0\007') - return - state = 1 if s['status'] == 'downloading' else 2 - self._downloader._send_console_code(f'\033]9;4;{state};{int(percent)}\007') + progress_dict), _ProgressState.from_dict(s), s.get('_progress')) def _retry_download(self, err, count, retries): # While this is not an extractor, it behaves similar to one and diff --git a/yt_dlp/utils/_utils.py b/yt_dlp/utils/_utils.py index 699bf1e7f..ef69eedd9 100644 --- a/yt_dlp/utils/_utils.py +++ b/yt_dlp/utils/_utils.py @@ -8,6 +8,7 @@ import datetime as dt import email.header import email.utils +import enum import errno import functools import hashlib @@ -5656,3 +5657,32 @@ def stdout(self, message): def stderr(self, message): if self._ydl: self._ydl.to_stderr(message) + + +class _ProgressState(enum.Enum): + """ + Represents a state for a progress bar. + + See: https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC + """ + + HIDDEN = 0 + INDETERMINATE = 3 + VISIBLE = 1 + WARNING = 4 + ERROR = 2 + + @classmethod + def from_dict(cls, s, /): + if s['status'] == 'finished': + return cls.INDETERMINATE + + # Not currently used + if s['status'] == 'error': + return cls.ERROR + + return cls.INDETERMINATE if s.get('_percent') is None else cls.VISIBLE + + def get_ansi_escape(self, /, percent=None): + percent = 0 if percent is None else int(percent) + return f'\033]9;4;{self.value};{percent}\007'