michael@0: # Text progress bar library, like curl or scp.
michael@0: 
michael@0: from datetime import datetime, timedelta
michael@0: import math
michael@0: import sys
michael@0: 
michael@0: if sys.platform.startswith('win'):
michael@0:     from terminal_win import Terminal
michael@0: else:
michael@0:     from terminal_unix import Terminal
michael@0: 
michael@0: class NullProgressBar(object):
michael@0:     def update(self, current, data): pass
michael@0:     def poke(self): pass
michael@0:     def finish(self, complete=True): pass
michael@0:     def beginline(self): pass
michael@0:     def message(self, msg): sys.stdout.write(msg + '\n')
michael@0:     def update_granularity(self): return timedelta.max
michael@0: 
michael@0: class ProgressBar(object):
michael@0:     def __init__(self, limit, fmt):
michael@0:         assert self.conservative_isatty()
michael@0: 
michael@0:         self.prior = None
michael@0:         self.atLineStart = True
michael@0:         self.counters_fmt = fmt # [{str:str}] Describtion of how to lay out each
michael@0:                                 #             field in the counters map.
michael@0:         self.limit = limit # int: The value of 'current' equal to 100%.
michael@0:         self.limit_digits = int(math.ceil(math.log10(self.limit))) # int: max digits in limit
michael@0:         self.t0 = datetime.now() # datetime: The start time.
michael@0: 
michael@0:         # Compute the width of the counters and build the format string.
michael@0:         self.counters_width = 1 # [
michael@0:         for layout in self.counters_fmt:
michael@0:             self.counters_width += self.limit_digits
michael@0:             self.counters_width += 1 # | (or ']' for the last one)
michael@0: 
michael@0:         self.barlen = 64 - self.counters_width
michael@0: 
michael@0:     def update_granularity(self):
michael@0:         return timedelta(seconds=0.1)
michael@0: 
michael@0:     def update(self, current, data):
michael@0:         # Record prior for poke.
michael@0:         self.prior = (current, data)
michael@0:         self.atLineStart = False
michael@0: 
michael@0:         # Build counters string.
michael@0:         sys.stdout.write('\r[')
michael@0:         for layout in self.counters_fmt:
michael@0:             Terminal.set_color(layout['color'])
michael@0:             sys.stdout.write(('%' + str(self.limit_digits) + 'd') % data[layout['value']])
michael@0:             Terminal.reset_color()
michael@0:             if layout != self.counters_fmt[-1]:
michael@0:                 sys.stdout.write('|')
michael@0:             else:
michael@0:                 sys.stdout.write('] ')
michael@0: 
michael@0:         # Build the bar.
michael@0:         pct = int(100.0 * current / self.limit)
michael@0:         sys.stdout.write('%3d%% ' % pct)
michael@0: 
michael@0:         barlen = int(1.0 * self.barlen * current / self.limit) - 1
michael@0:         bar = '=' * barlen + '>' + ' ' * (self.barlen - barlen - 1)
michael@0:         sys.stdout.write(bar + '|')
michael@0: 
michael@0:         # Update the bar.
michael@0:         dt = datetime.now() - self.t0
michael@0:         dt = dt.seconds + dt.microseconds * 1e-6
michael@0:         sys.stdout.write('%6.1fs' % dt)
michael@0:         Terminal.clear_right()
michael@0: 
michael@0:         # Force redisplay, since we didn't write a \n.
michael@0:         sys.stdout.flush()
michael@0: 
michael@0:     def poke(self):
michael@0:         if not self.prior:
michael@0:             return
michael@0:         self.update(*self.prior)
michael@0: 
michael@0:     def finish(self, complete=True):
michael@0:         final_count = self.limit if complete else self.prior[0]
michael@0:         self.update(final_count, self.prior[1])
michael@0:         sys.stdout.write('\n')
michael@0: 
michael@0:     def beginline(self):
michael@0:         if not self.atLineStart:
michael@0:             sys.stdout.write('\n')
michael@0:             self.atLineStart = True
michael@0: 
michael@0:     def message(self, msg):
michael@0:         self.beginline()
michael@0:         sys.stdout.write(msg)
michael@0:         sys.stdout.write('\n')
michael@0: 
michael@0:     @staticmethod
michael@0:     def conservative_isatty():
michael@0:         """
michael@0:         Prefer erring on the side of caution and not using terminal commands if
michael@0:         the current output stream may be a file.  We explicitly check for the
michael@0:         Android platform because terminal commands work poorly over ADB's
michael@0:         redirection.
michael@0:         """
michael@0:         try:
michael@0:             import android
michael@0:             return False
michael@0:         except ImportError:
michael@0:             return sys.stdout.isatty()
michael@0:         return False