xref: /aosp_15_r20/external/pigweed/pw_console/py/pw_console/progress_bar/progress_bar_task_counter.py (revision 61c4878ac05f98d0ceed94b57d316916de578985) 
1# Copyright 2021 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Container class for a single progress bar task."""
15
16from dataclasses import dataclass
17
18from prompt_toolkit.application import get_app_or_none
19from prompt_toolkit.shortcuts.progress_bar import ProgressBarCounter
20
21
22def _redraw_ui() -> None:
23    """Signal the prompt_toolkit app to re-draw"""
24    pt_app = get_app_or_none()
25    if pt_app:
26        pt_app.invalidate()
27
28
29@dataclass
30class ProgressBarTaskCounter:
31    """Class to hold a single progress bar state."""
32
33    name: str
34    total: int
35    count: int = 0
36    completed: bool = False
37    canceled: bool = False
38    prompt_toolkit_counter: ProgressBarCounter | None = None
39
40    def mark_canceled(self):
41        self.canceled = True
42        self.prompt_toolkit_counter.stopped = True  # type: ignore
43
44    def mark_completed(self):
45        self.completed = True
46        self.prompt_toolkit_counter.done = True  # type: ignore
47
48    def check_completion(self) -> None:
49        # Check for completion
50        if self.count >= self.total:
51            self.mark_completed()
52
53    def stop_updating_prompt_toolkit_counter(self) -> None:
54        """If count is over total, stop updating the prompt_toolkit ETA."""
55        if self.count >= self.total:
56            self.prompt_toolkit_counter.done = True  # type: ignore
57
58    def update(self, count: int = 1) -> None:
59        """Increment this counter."""
60        self.count += count
61
62        if self.prompt_toolkit_counter:
63            self.prompt_toolkit_counter.items_completed += count
64            self.stop_updating_prompt_toolkit_counter()
65            _redraw_ui()
66
67    def set_new_total(self, new_total: int) -> None:
68        """Set a new total count."""
69        self.count = new_total
70
71        if self.prompt_toolkit_counter:
72            self.prompt_toolkit_counter.items_completed = new_total
73            self.stop_updating_prompt_toolkit_counter()
74            _redraw_ui()
75