BmapCopy: implement progress inicator

I got a feature request from users to implement a progress bar in order to
show that the process is alive. This commit implements it.

I add a possibility to configure the BmapCopy class to print the progress
indicator to a user-defined file object with a user-defined pattern.

Change-Id: Ida0331a692163cbcf909ebf984917e2466eb0b70
Signed-off-by: Artem Bityutskiy <artem.bityutskiy@intel.com>

diff --git a/TODO b/TODO
index b40062f..b78b167 100644 (file)
--- a/TODO
+++ b/TODO
@@ -2,5 +2,3 @@
    signatures as documented here:
    http://www.garykessler.net/library/file_sigs.html
 2. Add a real man page to the project.
-3. Add a progress indicator: users want to see that the flashing is going on
-   while waiting, because it may take minutes.

diff --git a/bmaptool b/bmaptool
index 92d02c7..daa565f 100755 (executable)
--- a/bmaptool
+++ b/bmaptool
@@ -78,6 +78,9 @@ def copy_command(args, log):
         log.error(str(err))
         raise SystemExit(1)
 
+    # Print the progress indicator while copying
+    writer.set_progress_indicator(sys.stderr, "bmaptool: info: %d%% copied")
+
     start_time = time.time()
     if not args.bmap:
         log.warning("no bmap, flashing will be slow! (use --bmap option)")

diff --git a/bmaptools/BmapCopy.py b/bmaptools/BmapCopy.py
index 9722a1f..60eb048 100644 (file)
--- a/bmaptools/BmapCopy.py
+++ b/bmaptools/BmapCopy.py
@@ -101,7 +101,30 @@ class BmapCopy:
     not. To explicitly synchronize it, use the 'sync()' method.
 
     This class supports all the bmap format versions up version
-    'SUPPORTED_BMAP_VERSION'. """
+    'SUPPORTED_BMAP_VERSION'.
+
+    It is possible to have a simple progress indicator while copying the image.
+    Use the 'set_progress_indicator()' method. """
+
+    def set_progress_indicator(self, file_obj, format_string):
+        """ Setup the progress indicator which shows how much data has been
+        copied in percent.
+
+        The 'file_obj' argument is the console file object where the progress
+        has to be printed to. Pass 'None' to disable the progress indicator.
+
+        The 'format_string' argument is the format string for the progress
+        indicator. It has to contain a single '%d' placeholder which will be
+        substitutes with copied data in percent.
+
+        The progress indicator is not printed when copying a compressed image
+        without bmap, because the file size is unknown in this case. """
+
+        self._progress_file = file_obj
+        if format_string:
+            self._progress_format = format_string
+        else:
+            self._progress_format = "Copied %d%%"
 
     def _initialize_sizes(self, image_size):
         """ This function is only used when the there is no bmap. It
@@ -270,6 +293,11 @@ class BmapCopy:
         self._f_bmap = None
         self._f_bmap_path = None
 
+        self._progress_started = None
+        self._progress_file = None
+        self._progress_format = None
+        self.set_progress_indicator(None, None)
+
         if hasattr(dest, "write"):
             self._f_dest = dest
             self._dest_path = dest.name
@@ -323,6 +351,31 @@ class BmapCopy:
         if self._f_bmap_needs_close:
             self._f_bmap.close()
 
+    def _update_progress(self, blocks_written):
+        """ Print the progress indicator if the mapped area size is known and
+        if the indicator has been enabled by assigning a console file object to
+        the 'progress_file' attribute. """
+
+        if self._progress_file and self.mapped_cnt:
+            assert blocks_written <= self.mapped_cnt
+
+            percent = int((float(blocks_written) / self.mapped_cnt) * 100)
+
+            # This is a little trick we do in order to make sure that the next
+            # message will always start from a new line - we switch to the new
+            # line after each progress update and move the cursor up. As an
+            # example, this is useful when the copying is interrupted by an
+            # exception - the error message will start form new line.
+            if self._progress_started:
+                # The "move cursor up" escape sequence
+                self._progress_file.write('\033[1A')
+            else:
+                self._progress_started = True
+
+            fmt = '\r' + self._progress_format + '\n'
+            self._progress_file.write(fmt % percent)
+            self._progress_file.flush()
+
     def _get_block_ranges(self):
         """ This is a helper generator that parses the bmap XML file and for
         each block range in the XML file it yields ('first', 'last', 'sha1')
@@ -447,7 +500,7 @@ class BmapCopy:
         self._batch_queue.put(None)
 
     def copy(self, sync = True, verify = True):
-        """ Copy the image to the destination file using bmap. The sync
+        """ Copy the image to the destination file using bmap. The 'sync'
         argument defines whether the destination file has to be synchronized
         upon return.  The 'verify' argument defines whether the SHA1 checksum
         has to be verified while copying. """
@@ -466,6 +519,7 @@ class BmapCopy:
         blocks_written = 0
         bytes_written = 0
         fsync_last = 0
+        self._progress_started = False
 
         # Read the image in '_batch_blocks' chunks and write them to the
         # destination file
@@ -503,6 +557,8 @@ class BmapCopy:
             blocks_written += (end - start + 1)
             bytes_written += len(buf)
 
+            self._update_progress(blocks_written)
+
         if not self.image_size:
             # The image size was unknown up until now, probably because this is
             # a compressed image. Initialize the corresponding class attributes