| # Copyright 2021 The Pigweed Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| # use this file except in compliance with the License. You may obtain a copy of |
| # the License at |
| # |
| # https://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| # License for the specific language governing permissions and limitations under |
| # the License. |
| """Container class for a single progress bar task.""" |
| |
| from dataclasses import dataclass |
| from typing import Optional |
| |
| from prompt_toolkit.application import get_app_or_none |
| from prompt_toolkit.shortcuts.progress_bar import ProgressBarCounter |
| |
| |
| def _redraw_ui() -> None: |
| """Signal the prompt_toolkit app to re-draw""" |
| pt_app = get_app_or_none() |
| if pt_app: |
| pt_app.invalidate() |
| |
| |
| @dataclass |
| class ProgressBarTaskCounter: |
| """Class to hold a single progress bar state.""" |
| name: str |
| total: int |
| count: int = 0 |
| completed: bool = False |
| canceled: bool = False |
| prompt_toolkit_counter: Optional[ProgressBarCounter] = None |
| |
| def mark_canceled(self): |
| self.canceled = True |
| self.prompt_toolkit_counter.stopped = True # type: ignore |
| |
| def mark_completed(self): |
| self.completed = True |
| self.prompt_toolkit_counter.done = True # type: ignore |
| |
| def check_completion(self) -> None: |
| # Check for completion |
| if self.count >= self.total: |
| self.mark_completed() |
| |
| def stop_updating_prompt_toolkit_counter(self) -> None: |
| """If count is over total, stop updating the prompt_toolkit ETA.""" |
| if self.count >= self.total: |
| self.prompt_toolkit_counter.done = True # type: ignore |
| |
| def update(self, count: int = 1) -> None: |
| """Increment this counter.""" |
| self.count += count |
| |
| if self.prompt_toolkit_counter: |
| self.prompt_toolkit_counter.items_completed += count |
| self.stop_updating_prompt_toolkit_counter() |
| _redraw_ui() |
| |
| def set_new_total(self, new_total: int) -> None: |
| """Set a new total count.""" |
| self.count = new_total |
| |
| if self.prompt_toolkit_counter: |
| self.prompt_toolkit_counter.items_completed = new_total |
| self.stop_updating_prompt_toolkit_counter() |
| _redraw_ui() |
