• 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
34    name: str
35    total: int
36    count: int = 0
37    completed: bool = False
38    canceled: bool = False
39    prompt_toolkit_counter: Optional[ProgressBarCounter] = None
40
41    def mark_canceled(self):
42        self.canceled = True
43        self.prompt_toolkit_counter.stopped = True  # type: ignore
44
45    def mark_completed(self):
46        self.completed = True
47        self.prompt_toolkit_counter.done = True  # type: ignore
48
49    def check_completion(self) -> None:
50        # Check for completion
51        if self.count >= self.total:
52            self.mark_completed()
53
54    def stop_updating_prompt_toolkit_counter(self) -> None:
55        """If count is over total, stop updating the prompt_toolkit ETA."""
56        if self.count >= self.total:
57            self.prompt_toolkit_counter.done = True  # type: ignore
58
59    def update(self, count: int = 1) -> None:
60        """Increment this counter."""
61        self.count += count
62
63        if self.prompt_toolkit_counter:
64            self.prompt_toolkit_counter.items_completed += count
65            self.stop_updating_prompt_toolkit_counter()
66            _redraw_ui()
67
68    def set_new_total(self, new_total: int) -> None:
69        """Set a new total count."""
70        self.count = new_total
71
72        if self.prompt_toolkit_counter:
73            self.prompt_toolkit_counter.items_completed = new_total
74            self.stop_updating_prompt_toolkit_counter()
75            _redraw_ui()
76