From 94865ed8ab561fa8f1a7aae9e1663fd7bea2b370 Mon Sep 17 00:00:00 2001
From: Yves Goergen <git@ygoe.de>
Date: Sun, 26 Apr 2020 15:48:00 +0200
Subject: [PATCH] Added optional progress output for writing/reading flash

---
 pyupdi.py   | 14 ++++++++++++--
 updi/nvm.py | 10 ++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/pyupdi.py b/pyupdi.py
index 63ed86b..97f194c 100755
--- a/pyupdi.py
+++ b/pyupdi.py
@@ -78,6 +78,8 @@ def _main():
                         help="Read out the fuse-bits")
     parser.add_argument("-v", "--verbose", action="store_true",
                         help="Set verbose mode")
+    parser.add_argument("-p", "--progress", action="store_true",
+                        help="Show progress output")
 
     args = parser.parse_args(sys.argv[1:])
 
@@ -95,6 +97,8 @@ def _main():
     nvm = UpdiNvmProgrammer(comport=args.comport,
                             baud=args.baudrate,
                             device=Device(args.device))
+    if args.progress:
+        nvm.enable_progress()
     if not args.reset: # any action except reset
         # Reteieve info before building the stack to be sure its the correct device
         nvm.get_device_info()
@@ -131,22 +135,28 @@ def _process(nvm, args):
                 if not _set_fuse(nvm, fusenum, value):
                     return False
     if args.flash is not None:
-        return _flash_file(nvm, args.flash)
+        return _flash_file(nvm, args.flash, args.progress)
     if args.readfuses:
         if not _read_fuses(nvm):
             return False
     return True
 
 
-def _flash_file(nvm, filename):
+def _flash_file(nvm, filename, progress):
     data, start_address = nvm.load_ihex(filename)
 
     fail = False
 
+    if progress:
+        print("Erasing chip...")
     nvm.chip_erase()
+    if progress:
+        print("Writing flash...")
     nvm.write_flash(start_address, data)
 
     # Read out again
+    if progress:
+        print("Verifying flash...")
     readback = nvm.read_flash(nvm.device.flash_start, len(data))
     for i, _ in enumerate(data):
         if data[i] != readback[i]:
diff --git a/updi/nvm.py b/updi/nvm.py
index 6b61e2d..a1ee6bd 100644
--- a/updi/nvm.py
+++ b/updi/nvm.py
@@ -18,6 +18,10 @@ def __init__(self, comport, baud, device):
         self.device = device
         self.progmode = False
         self.logger = logging.getLogger("nvm")
+        self.progress = False
+
+    def enable_progress(self):
+        self.progress = True
 
     def get_device_info(self):
         """
@@ -79,8 +83,11 @@ def read_flash(self, address, size):
             raise Exception("Only full page aligned flash supported.")
 
         data = []
+        start_address = address
         # Read out page-wise for convenience
         for _ in range(pages):
+            if self.progress:
+                print("  {0:3.0f}%".format((address - start_address) * 100 / size), end = "\r")
             self.logger.info("Reading page at 0x{0:04X}".format(address))
             data += (self.application.read_data_words(address, self.device.flash_pagesize >> 1))
             address += self.device.flash_pagesize
@@ -101,7 +108,10 @@ def write_flash(self, address, data):
         pages = self.page_data(data, self.device.flash_pagesize)
 
         # Program each page
+        start_address = address
         for page in pages:
+            if self.progress:
+                print("  {0:3.0f}%".format((address - start_address) * 100 / len(data)), end = "\r")
             self.logger.info("Writing page at 0x{0:04X}".format(address))
             self.application.write_nvm(address, page)
             address += len(page)