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