• 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"""Pigweed Console ProgressBar implementation.
15
16Designed to be embedded in an existing prompt_toolkit full screen
17application."""
18
19import functools
20from typing import (
21    Iterable,
22    List,
23    Optional,
24    Sequence,
25)
26
27from prompt_toolkit.filters import Condition
28from prompt_toolkit.formatted_text import AnyFormattedText
29from prompt_toolkit.layout import (
30    ConditionalContainer,
31    FormattedTextControl,
32    HSplit,
33    VSplit,
34    Window,
35)
36from prompt_toolkit.layout.dimension import AnyDimension, D
37from prompt_toolkit.styles import BaseStyle
38
39from prompt_toolkit.shortcuts.progress_bar import (
40    ProgressBar,
41    ProgressBarCounter,
42)
43from prompt_toolkit.shortcuts.progress_bar.base import _ProgressControl
44from prompt_toolkit.shortcuts.progress_bar.formatters import (
45    Formatter,
46    IterationsPerSecond,
47    Text,
48    TimeLeft,
49    create_default_formatters,
50)
51
52
53class TextIfNotHidden(Text):
54    def format(
55        self,
56        progress_bar: ProgressBar,
57        progress: 'ProgressBarCounter[object]',
58        width: int,
59    ) -> AnyFormattedText:
60        formatted_text = super().format(progress_bar, progress, width)
61        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
62
63            formatted_text = [('', ' ' * width)]
64        return formatted_text
65
66
67class IterationsPerSecondIfNotHidden(IterationsPerSecond):
68    def format(
69        self,
70        progress_bar: ProgressBar,
71        progress: 'ProgressBarCounter[object]',
72        width: int,
73    ) -> AnyFormattedText:
74        formatted_text = super().format(progress_bar, progress, width)
75        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
76            formatted_text = [('class:iterations-per-second', ' ' * width)]
77        return formatted_text
78
79
80class TimeLeftIfNotHidden(TimeLeft):
81    def format(
82        self,
83        progress_bar: ProgressBar,
84        progress: 'ProgressBarCounter[object]',
85        width: int,
86    ) -> AnyFormattedText:
87        formatted_text = super().format(progress_bar, progress, width)
88        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
89            formatted_text = [('class:time-left', ' ' * width)]
90        return formatted_text
91
92
93class ProgressBarImpl:
94    """ProgressBar for rendering in an existing prompt_toolkit application."""
95    def __init__(
96        self,
97        title: AnyFormattedText = None,
98        formatters: Optional[Sequence[Formatter]] = None,
99        style: Optional[BaseStyle] = None,
100    ) -> None:
101
102        self.title = title
103        self.formatters = formatters or create_default_formatters()
104        self.counters: List[ProgressBarCounter[object]] = []
105        self.style = style
106
107        # Create UI Application.
108        title_toolbar = ConditionalContainer(
109            Window(
110                FormattedTextControl(lambda: self.title),
111                height=1,
112                style='class:progressbar,title',
113            ),
114            filter=Condition(lambda: self.title is not None),
115        )
116
117        def width_for_formatter(formatter: Formatter) -> AnyDimension:
118            # Needs to be passed as callable (partial) to the 'width'
119            # parameter, because we want to call it on every resize.
120            return formatter.get_width(progress_bar=self)  # type: ignore
121
122        progress_controls = [
123            Window(
124                content=_ProgressControl(self, f),  # type: ignore
125                width=functools.partial(width_for_formatter, f),
126            ) for f in self.formatters
127        ]
128
129        self.container = HSplit([
130            title_toolbar,
131            VSplit(
132                progress_controls,
133                height=lambda: D(min=len(self.counters),
134                                 max=len(self.counters)),
135            ),
136        ])
137
138    def __pt_container__(self):
139        return self.container
140
141    def __exit__(self, *a: object) -> None:
142        pass
143
144    def __call__(
145        self,
146        data: Optional[Iterable] = None,
147        label: AnyFormattedText = '',
148        remove_when_done: bool = False,
149        total: Optional[int] = None,
150    ) -> 'ProgressBarCounter':
151        """
152        Start a new counter.
153
154        :param label: Title text or description for this progress. (This can be
155            formatted text as well).
156        :param remove_when_done: When `True`, hide this progress bar.
157        :param total: Specify the maximum value if it can't be calculated by
158            calling ``len``.
159        """
160        counter = ProgressBarCounter(
161            self,  # type: ignore
162            data,
163            label=label,
164            remove_when_done=remove_when_done,
165            total=total)
166        self.counters.append(counter)
167        return counter
168