Coverage for src / mafw / ui / rich_user_interface.py: 100%
67 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-09 09:08 +0000
1"""
2The rich user interface.
4The module provides an implementation of the abstract user interface that takes advantage from the `rich` library.
5Progress bars and spinners are shown during the processor execution along with log messages including markup language.
6In order for this logging message to appear properly rendered, the logger should be connected to a RichHandler.
7"""
9# Copyright 2025 European Union
10# Author: Bulgheroni Antonio (antonio.bulgheroni@ec.europa.eu)
11# SPDX-License-Identifier: EUPL-1.2
12import logging
13from contextlib import contextmanager
14from types import TracebackType
15from typing import Any, Generator, Self
17import rich.prompt
18from rich.progress import Progress, SpinnerColumn, TaskID, TimeElapsedColumn
20from mafw.enumerators import ProcessorStatus
21from mafw.ui.abstract_user_interface import UserInterfaceBase
23log = logging.getLogger(__name__)
26class RichInterface(UserInterfaceBase):
27 """
28 Implementation of the interface for rich.
30 :param progress_kws: A dictionary of keywords passed to the `rich.Progress`. Defaults to None
31 :type progress_kws: dict, Optional
32 """
34 name = 'rich'
36 def __init__(self, progress_kws: dict[str, Any] | None = None) -> None:
37 if progress_kws is None:
38 progress_kws = dict(auto_refresh=True, expand=True)
40 self.progress = Progress(SpinnerColumn(), *Progress.get_default_columns(), TimeElapsedColumn(), **progress_kws)
41 self.task_dict: dict[str, TaskID] = {}
43 def __enter__(self) -> Self:
44 """
45 Context enter dunder.
47 It manually starts the progress extension and then return the class instance.
48 """
49 self.progress.start()
50 return self
52 def __exit__(
53 self, type_: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None
54 ) -> None:
55 """
56 Context exit dunder.
58 It manually stops the progress bar.
60 :param type_: Exception type.
61 :param value: Exception value.
62 :param traceback: Exception trace back.
63 """
64 self.progress.stop()
66 def create_task(
67 self,
68 task_name: str,
69 task_description: str = '',
70 completed: int = 0,
71 increment: int | None = None,
72 total: int | None = None,
73 **kwargs: Any,
74 ) -> None:
75 """
76 Create a new task.
78 :param task_name: A unique identifier for the task. You cannot have more than 1 task with the same name in
79 the whole execution. If you want to use the processor name, it is recommended to use the
80 :attr:`~mafw.processor.Processor.unique_name`.
81 :type task_name: str
82 :param task_description: A short description for the task. Defaults to ''.
83 :type task_description: str, Optional
84 :param completed: The amount of task already completed. Defaults to 0.
85 :type completed: int, Optional
86 :param increment: How much of the task has been done since last update. Defaults to None.
87 :type increment: int, Optional
88 :param total: The total amount of task. Defaults to None.
89 :type total: int, Optional
90 """
91 if task_name in self.task_dict:
92 log.warning('A task with this name (%s) already exists. Replacing it with the new one.' % task_name)
93 log.warning('Be sure to use unique names.')
95 self.task_dict[task_name] = self.progress.add_task(task_description, total=total, completed=completed)
97 def update_task(
98 self,
99 task_name: str,
100 completed: int | None = None,
101 increment: int | None = None,
102 total: int | None = None,
103 **kwargs: Any,
104 ) -> None:
105 """
106 Update an existing task.
108 :param task_name: A unique identifier for the task. You cannot have more than one task with the same name in
109 the whole execution. If you want to use the processor name, it is recommended to use the
110 :attr:`~~mafw.processor.Processor.replica_name`.
111 :type task_name: str
112 :param completed: The amount of task already completed. Defaults to 0.
113 :type completed: int, Optional
114 :param increment: How much of the task has been done since last update. Defaults to None.
115 :type increment: int, Optional
116 :param total: The total amount of task. Defaults to None.
117 :type total: int, Optional
118 """
119 if task_name not in self.task_dict:
120 log.warning('A task with this name (%s) does not exist.' % task_name)
121 log.warning('Skipping updates')
122 return
124 if completed is None and total is None:
125 visible = True
126 else:
127 visible = completed != total
129 self.progress.update(
130 self.task_dict[task_name], completed=completed, advance=increment, total=total, visible=visible
131 )
133 def display_progress_message(self, message: str, i_item: int, n_item: int | None, frequency: float) -> None:
134 if self._is_time_to_display_lopping_message(i_item, n_item, frequency):
135 if n_item is None:
136 n_item = max(1000, i_item)
137 width = len(str(n_item))
138 counter = f'[{i_item + 1:>{width}}/{n_item}] '
139 msg = counter + message
140 log.info(msg)
142 def change_of_processor_status(
143 self, processor_name: str, old_status: ProcessorStatus, new_status: ProcessorStatus
144 ) -> None:
145 """
146 Display a message when a processor status changes.
148 This method logs a debug message indicating that a processor has changed its status.
149 The message uses rich markup to highlight the processor name and new status.
151 :param processor_name: The name of the processor whose status has changed.
152 :type processor_name: str
153 :param old_status: The previous status of the processor.
154 :type old_status: ProcessorStatus
155 :param new_status: The new status of the processor.
156 :type new_status: ProcessorStatus
157 """
158 msg = f'[red]{processor_name}[/red] is [bold]{new_status}[/bold]'
159 log.debug(msg)
161 @contextmanager
162 def enter_interactive_mode(self) -> Generator[None, Any, None]:
163 """
164 Context manager to temporarily switch to interactive mode.
166 This method temporarily stops the progress display to allow for interactive input
167 while preserving the original transient state. After yielding control, it restores
168 the progress display with appropriate spacing to avoid overwriting previous output.
170 .. versionadded:: v2.0.0
172 .. note::
173 This method should be used within a ``with`` statement to ensure proper cleanup.
174 """
175 transient = self.progress.live.transient # save the old value
176 self.progress.live.transient = True
177 self.progress.stop()
178 self.progress.live.transient = transient # restore the old value
179 try:
180 yield
181 finally:
182 # make space for the progress to use so it doesn't overwrite any previous lines
183 visible_tasks = [task for task in self.progress.tasks if task.visible]
184 print('\n' * (len(visible_tasks) - 2))
185 self.progress.start()
187 def prompt_question(self, question: str, **kwargs: Any) -> Any:
188 """
189 Prompt the user with a question and return their response.
191 This method uses the rich library's prompt functionality to ask the user a question.
192 It supports various prompt types including confirmation, input, and choice prompts.
194 .. versionadded:: v2.0.0
196 :param question: The question to ask the user.
197 :type question: str
198 :param kwargs: Additional arguments to pass to the prompt function.
199 :return: The user's response based on the prompt type.
200 :rtype: Any
201 :param prompt_type: The type of prompt to use. Defaults to :class:`rich.prompt.Confirm`.
202 :param console: The console to use for the prompt. Defaults to None.
203 :param password: Whether to hide input when prompting for passwords. Defaults to False.
204 :param choices: List of valid choices for choice prompts. Defaults to None.
205 :param default: Default value for prompts that support it. Defaults to None.
206 :param show_default: Whether to show the default value. Defaults to True.
207 :param show_choices: Whether to show available choices. Defaults to True.
208 :param case_sensitive: Whether choices are case sensitive. Defaults to True.
209 """
210 prompt_type = kwargs.pop('prompt_type', rich.prompt.Confirm)
211 console = kwargs.pop('console', None)
212 password = kwargs.pop('password', False)
213 choices = kwargs.pop('choices', None)
214 default = kwargs.pop('default', None)
215 show_default = kwargs.pop('show_default', True)
216 show_choices = kwargs.pop('show_choices', True)
217 case_sensitive = kwargs.pop('case_sensitive', True)
219 return prompt_type.ask(
220 question,
221 console=console,
222 password=password,
223 choices=choices,
224 default=default,
225 case_sensitive=case_sensitive,
226 show_default=show_default,
227 show_choices=show_choices,
228 )