pw_console/py/pw_console/progress_bar/__init__.py

# 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.
"""Pigweed Console progress bar functions."""
from pw_console.progress_bar.progress_bar_state import TASKS_CONTEXTVAR
from pw_console.progress_bar.progress_bar_task_counter import (
    ProgressBarTaskCounter)

__all__ = [
    'start_progress',
    'update_progress',
]


def start_progress(task_name: str, total: int, hide_eta=False):
    progress_state = TASKS_CONTEXTVAR.get()

    progress_state.startup_progress_bar_impl()

    assert progress_state.instance is not None

    progress_state.tasks[task_name] = ProgressBarTaskCounter(
        name=task_name,
        total=total,
        prompt_toolkit_counter=progress_state.instance(range(total),
                                                       label=task_name))
    ptc = progress_state.tasks[task_name].prompt_toolkit_counter
    ptc.hide_eta = hide_eta  # type: ignore


def update_progress(task_name: str,
                    count=1,
                    completed=False,
                    canceled=False,
                    new_total=None):
    progress_state = TASKS_CONTEXTVAR.get()
    # The caller may not actually get canceled and will continue trying to
    # update after an interrupt.
    if task_name not in progress_state.tasks:
        return

    # Take one action
    if completed:
        progress_state.tasks[task_name].mark_completed()
    elif canceled:
        progress_state.tasks[task_name].mark_canceled()
    elif new_total:
        progress_state.tasks[task_name].set_new_total(new_total)
    else:
        progress_state.tasks[task_name].update(count)

    # Check if all tasks are complete
    if (progress_state.instance is not None
            and progress_state.all_tasks_complete):
        if hasattr(progress_state.instance, '__exit__'):
            progress_state.instance.__exit__()