pw_console/py/pw_console/progress_bar/progress_bar_impl.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 ProgressBar implementation.

Designed to be embedded in an existing prompt_toolkit full screen
application."""

import functools
from typing import (
    Iterable,
    List,
    Optional,
    Sequence,
)

from prompt_toolkit.filters import Condition
from prompt_toolkit.formatted_text import AnyFormattedText
from prompt_toolkit.layout import (
    ConditionalContainer,
    FormattedTextControl,
    HSplit,
    VSplit,
    Window,
)
from prompt_toolkit.layout.dimension import AnyDimension, D
from prompt_toolkit.styles import BaseStyle

from prompt_toolkit.shortcuts.progress_bar import (
    ProgressBar,
    ProgressBarCounter,
)
from prompt_toolkit.shortcuts.progress_bar.base import _ProgressControl
from prompt_toolkit.shortcuts.progress_bar.formatters import (
    Formatter,
    IterationsPerSecond,
    Text,
    TimeLeft,
    create_default_formatters,
)


class TextIfNotHidden(Text):
    def format(
        self,
        progress_bar: ProgressBar,
        progress: 'ProgressBarCounter[object]',
        width: int,
    ) -> AnyFormattedText:
        formatted_text = super().format(progress_bar, progress, width)
        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore

            formatted_text = [('', ' ' * width)]
        return formatted_text


class IterationsPerSecondIfNotHidden(IterationsPerSecond):
    def format(
        self,
        progress_bar: ProgressBar,
        progress: 'ProgressBarCounter[object]',
        width: int,
    ) -> AnyFormattedText:
        formatted_text = super().format(progress_bar, progress, width)
        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
            formatted_text = [('class:iterations-per-second', ' ' * width)]
        return formatted_text


class TimeLeftIfNotHidden(TimeLeft):
    def format(
        self,
        progress_bar: ProgressBar,
        progress: 'ProgressBarCounter[object]',
        width: int,
    ) -> AnyFormattedText:
        formatted_text = super().format(progress_bar, progress, width)
        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
            formatted_text = [('class:time-left', ' ' * width)]
        return formatted_text


class ProgressBarImpl:
    """ProgressBar for rendering in an existing prompt_toolkit application."""
    def __init__(
        self,
        title: AnyFormattedText = None,
        formatters: Optional[Sequence[Formatter]] = None,
        style: Optional[BaseStyle] = None,
    ) -> None:

        self.title = title
        self.formatters = formatters or create_default_formatters()
        self.counters: List[ProgressBarCounter[object]] = []
        self.style = style

        # Create UI Application.
        title_toolbar = ConditionalContainer(
            Window(
                FormattedTextControl(lambda: self.title),
                height=1,
                style='class:progressbar,title',
            ),
            filter=Condition(lambda: self.title is not None),
        )

        def width_for_formatter(formatter: Formatter) -> AnyDimension:
            # Needs to be passed as callable (partial) to the 'width'
            # parameter, because we want to call it on every resize.
            return formatter.get_width(progress_bar=self)  # type: ignore

        progress_controls = [
            Window(
                content=_ProgressControl(self, f),  # type: ignore
                width=functools.partial(width_for_formatter, f),
            ) for f in self.formatters
        ]

        self.container = HSplit([
            title_toolbar,
            VSplit(
                progress_controls,
                height=lambda: D(min=len(self.counters),
                                 max=len(self.counters)),
            ),
        ])

    def __pt_container__(self):
        return self.container

    def __exit__(self, *a: object) -> None:
        pass

    def __call__(
        self,
        data: Optional[Iterable] = None,
        label: AnyFormattedText = '',
        remove_when_done: bool = False,
        total: Optional[int] = None,
    ) -> 'ProgressBarCounter':
        """
        Start a new counter.

        :param label: Title text or description for this progress. (This can be
            formatted text as well).
        :param remove_when_done: When `True`, hide this progress bar.
        :param total: Specify the maximum value if it can't be calculated by
            calling ``len``.
        """
        counter = ProgressBarCounter(
            self,  # type: ignore
            data,
            label=label,
            remove_when_done=remove_when_done,
            total=total)
        self.counters.append(counter)
        return counter