1# Copyright 2013 OpenStack Foundation
2# All Rights Reserved.
3#
4#    Licensed under the Apache License, Version 2.0 (the "License"); you may
5#    not use this file except in compliance with the License. You may obtain
6#    a copy of the License at
7#
8#         http://www.apache.org/licenses/LICENSE-2.0
9#
10#    Unless required by applicable law or agreed to in writing, software
11#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13#    License for the specific language governing permissions and limitations
14#    under the License.
15
16import sys
17
18
19class _ProgressBarBase(object):
20    """A progress bar provider for a wrapped obect.
21
22    Base abstract class used by specific class wrapper to show
23    a progress bar when the wrapped object are consumed.
24
25    :param wrapped: Object to wrap that hold data to be consumed.
26    :param totalsize: The total size of the data in the wrapped object.
27
28    :note: The progress will be displayed only if sys.stdout is a tty.
29    """
30
31    def __init__(self, wrapped, totalsize):
32        self._wrapped = wrapped
33        self._totalsize = float(totalsize)
34        self._show_progress = sys.stdout.isatty() and self._totalsize != 0
35        self._percent = 0
36
37    def _display_progress_bar(self, size_read):
38        if self._show_progress:
39            self._percent += size_read / self._totalsize
40            # Output something like this: [==========>             ] 49%
41            sys.stdout.write('\r[{0:<30}] {1:.0%}'.format(
42                '=' * int(round(self._percent * 29)) + '>', self._percent
43            ))
44            sys.stdout.flush()
45
46    def __getattr__(self, attr):
47        # Forward other attribute access to the wrapped object.
48        return getattr(self._wrapped, attr)
49
50
51class VerboseFileWrapper(_ProgressBarBase):
52    """A file wrapper with a progress bar.
53
54    The file wrapper shows and advances a progress bar whenever the
55    wrapped file's read method is called.
56    """
57
58    def read(self, *args, **kwargs):
59        data = self._wrapped.read(*args, **kwargs)
60        if data:
61            self._display_progress_bar(len(data))
62        else:
63            if self._show_progress:
64                # Break to a new line from the progress bar for incoming
65                # output.
66                sys.stdout.write('\n')
67        return data
68
69
70class VerboseIteratorWrapper(_ProgressBarBase):
71    """An iterator wrapper with a progress bar.
72
73    The iterator wrapper shows and advances a progress bar whenever the
74    wrapped data is consumed from the iterator.
75
76    :note: Use only with iterator that yield strings.
77    """
78
79    def __iter__(self):
80        return self
81
82    def next(self):
83        try:
84            data = next(self._wrapped)
85            # NOTE(mouad): Assuming that data is a string b/c otherwise calling
86            # len function will not make any sense.
87            self._display_progress_bar(len(data))
88            return data
89        except StopIteration:
90            if self._show_progress:
91                # Break to a new line from the progress bar for incoming
92                # output.
93                sys.stdout.write('\n')
94            raise
95
96    # In Python 3, __next__() has replaced next().
97    __next__ = next
98