• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
17from typing import Optional
18
19from prompt_toolkit.application import get_app_or_none
20from prompt_toolkit.shortcuts.progress_bar import ProgressBarCounter
21
22
23def _redraw_ui() -> None:
24    """Signal the prompt_toolkit app to re-draw"""
25    pt_app = get_app_or_none()
26    if pt_app:
27        pt_app.invalidate()
28
29
30@dataclass
31class ProgressBarTaskCounter:
32    """Class to hold a single progress bar state."""
33    name: str
34    total: int
35    count: int = 0
36    completed: bool = False
37    canceled: bool = False
38    prompt_toolkit_counter: Optional[ProgressBarCounter] = 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