Make Progress support an arbitrary number of columns.

tools/isolateserver_load_test.py

 #!/usr/bin/env python
 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Uploads a ton of stuff to isolateserver to test its handling.
 
 Generates an histogram with the latencies to download a just uploaded file.
 
 Note that it only looks at uploading and downloading and do not test
@@ skipping 34 matching lines @@
   def gen(self, size):
     """Returns a str containing random data from the pool of size |size|."""
     chunks = int(size / 1024)
     rest = size - (chunks*1024)
     data = ''.join(random.choice(self.pool) for _ in xrange(chunks))
     data += random.choice(self.pool)[:rest]
     return data
 
 
 class Progress(threading_utils.Progress):
-  def __init__(self, *args, **kwargs):
-    super(Progress, self).__init__(*args, **kwargs)
-    self.total = 0
-
-  def increment_index(self, name):
-    self.update_item(name, index=1)
-
-  def increment_count(self):
-    self.update_item('', size=1)
-
-  def gen_line(self, name):
-    """Generates the line to be printed.
-
-    |name| is actually the item size.
-    """
-    if name:
-      self.total += name
-      name = graph.to_units(name)
-
-    next_line = ('[%*d/%d] %6.2fs %8s %8s') % (
-        len(str(self.size)), self.index,
-        self.size,
-        time.time() - self.start,
-        name,
-        graph.to_units(self.total))
-    # Fill it with whitespace only if self.use_cr_only is set.
-    prefix = ''
-    if self.use_cr_only and self.last_printed_line:
-      prefix = '\r'
-    if self.use_cr_only:
-      suffix = ' ' * max(0, len(self.last_printed_line) - len(next_line))
-    else:
-      suffix = '\n'
-    return '%s%s%s' % (prefix, next_line, suffix), next_line
+  def _render_columns(self):
+    """Prints the size data as 'units'."""
+    columns_as_str = [
+        str(self._columns[0]),
+        graph.to_units(self._columns[1]).rjust(6),
+        str(self._columns[2]),
+    ]
+    max_len = max((len(columns_as_str[0]), len(columns_as_str[2])))
+    return '/'.join(i.rjust(max_len) for i in columns_as_str)
 
 
 def print_results(results, columns, buckets):
   delays = [i[0] for i in results if isinstance(i[0], float)]
   failures = [i for i in results if not isinstance(i[0], float)]
   sizes = [i[1] for i in results]
 
   print('%sSIZES%s (bytes):' % (colorama.Fore.RED, colorama.Fore.RESET))
   graph.print_histogram(
       graph.generate_histogram(sizes, buckets), columns, '%d')
@@ skipping 48 matching lines @@
       logging.info('download')
       # Only count download time when not in dry run time.
       # TODO(maruel): Count the number of retries!
       start = time.time()
       assert content == ''.join(api.fetch(hash_value, size))
     else:
       time.sleep(size / 10.)
     duration = max(0, time.time() - start)
   except isolateserver.MappingError as e:
     duration = str(e)
-  progress.increment_index(size if isinstance(duration, float) else 0)
+  if isinstance(duration, float):
+    progress.update_item('', index=1, data=size)
+  else:
+    progress.update_item('', index=1)
   return (duration, size)
 
 
 def main():
   colorama.init()
 
   parser = optparse.OptionParser(description=sys.modules[__name__].__doc__)
   parser.add_option(
       '-I', '--isolate-server',
       metavar='URL', default='',
@@ skipping 44 matching lines @@
       ' - Using %d thread,  items=%d,  max-size=%d,  mid-size=%d' % (
       options.threads, options.items, options.max_size, options.mid_size))
   if options.dry_run:
     print(' - %sDRY RUN MODE%s' % (colorama.Fore.GREEN, colorama.Fore.RESET))
 
   start = time.time()
 
   random_pool = Randomness()
   print(' - Generated pool after %.1fs' % (time.time() - start))
 
-  progress = Progress(options.items)
+  columns = [('index', 0), ('data', 0), ('size', options.items)]
+  progress = Progress(columns)
   api = isolateserver.get_storage_api(options.isolate_server, options.namespace)
   do_item = functools.partial(
       send_and_receive,
       random_pool,
       options.dry_run,
       api,
       progress)
 
   # TODO(maruel): Handle Ctrl-C should:
   # - Stop adding tasks.
   # - Stop scheduling tasks in ThreadPool.
   # - Wait for the remaining ungoing tasks to complete.
   # - Still print details and write the json file.
   with threading_utils.ThreadPoolWithProgress(
       progress, options.threads, options.threads, 0) as pool:
     if options.items:
       for _ in xrange(options.items):
         pool.add_task(0, do_item, gen_size(options.mid_size))
+        progress.print_update()
     elif options.max_size:
       # This one is approximate.
       total = 0
       while True:
         size = gen_size(options.mid_size)
-        progress.increment_count()
+        progress.update_item('', size=1)
         progress.print_update()
         pool.add_task(0, do_item, size)
         total += size
         if total >= options.max_size:
           break
     results = sorted(pool.join())
 
   print('')
   print(' - Took %.1fs.' % (time.time() - start))
   print('')
   print_results(results, options.columns, options.buckets)
   if options.dump:
     with open(options.dump, 'w') as f:
       json.dump(results, f, separators=(',',':'))
   return 0
 
 
 if __name__ == '__main__':
   sys.exit(main())