• 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            formatted_text = [('', ' ' * width)]
63        return formatted_text
64
65
66class IterationsPerSecondIfNotHidden(IterationsPerSecond):
67    def format(
68        self,
69        progress_bar: ProgressBar,
70        progress: 'ProgressBarCounter[object]',
71        width: int,
72    ) -> AnyFormattedText:
73        formatted_text = super().format(progress_bar, progress, width)
74        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
75            formatted_text = [('class:iterations-per-second', ' ' * width)]
76        return formatted_text
77
78
79class TimeLeftIfNotHidden(TimeLeft):
80    def format(
81        self,
82        progress_bar: ProgressBar,
83        progress: 'ProgressBarCounter[object]',
84        width: int,
85    ) -> AnyFormattedText:
86        formatted_text = super().format(progress_bar, progress, width)
87        if hasattr(progress, 'hide_eta') and progress.hide_eta:  # type: ignore
88            formatted_text = [('class:time-left', ' ' * width)]
89        return formatted_text
90
91
92class ProgressBarImpl:
93    """ProgressBar for rendering in an existing prompt_toolkit application."""
94
95    def __init__(
96        self,
97        title: AnyFormattedText = None,
98        formatters: Optional[Sequence[Formatter]] = None,
99        style: Optional[BaseStyle] = None,
100    ) -> None:
101        self.title = title
102        self.formatters = formatters or create_default_formatters()
103        self.counters: List[ProgressBarCounter[object]] = []
104        self.style = style
105
106        # Create UI Application.
107        title_toolbar = ConditionalContainer(
108            Window(
109                FormattedTextControl(lambda: self.title),
110                height=1,
111                style='class:progressbar,title',
112            ),
113            filter=Condition(lambda: self.title is not None),
114        )
115
116        def width_for_formatter(formatter: Formatter) -> AnyDimension:
117            # Needs to be passed as callable (partial) to the 'width'
118            # parameter, because we want to call it on every resize.
119            return formatter.get_width(progress_bar=self)  # type: ignore
120
121        progress_controls = [
122            Window(
123                content=_ProgressControl(self, f, None),  # type: ignore
124                width=functools.partial(width_for_formatter, f),
125            )
126            for f in self.formatters
127        ]
128
129        self.container = HSplit(
130            [
131                title_toolbar,
132                VSplit(
133                    progress_controls,
134                    height=lambda: D(
135                        min=len(self.counters), max=len(self.counters)
136                    ),
137                ),
138            ]
139        )
140
141    def __pt_container__(self):
142        return self.container
143
144    def __exit__(self, *a: object) -> None:
145        pass
146
147    def __call__(
148        self,
149        data: Optional[Iterable] = None,
150        label: AnyFormattedText = '',
151        remove_when_done: bool = False,
152        total: Optional[int] = None,
153    ) -> 'ProgressBarCounter':
154        """
155        Start a new counter.
156
157        :param label: Title text or description for this progress. (This can be
158            formatted text as well).
159        :param remove_when_done: When `True`, hide this progress bar.
160        :param total: Specify the maximum value if it can't be calculated by
161            calling ``len``.
162        """
163        counter = ProgressBarCounter(
164            self,  # type: ignore
165            data,
166            label=label,
167            remove_when_done=remove_when_done,
168            total=total,
169        )
170        self.counters.append(counter)
171        return counter
172