From 56978c3dde3440367438a89353ba73b32c8cdbc2 Mon Sep 17 00:00:00 2001
From: Rob Knegjens <rob@knegjens.net>
Date: Tue, 5 Apr 2022 16:05:28 -0700
Subject: [PATCH] tools/mpremote: Show progress indicator when copying large
 files.

When copying large files (> 2048 bytes) to or from a device with
`mpremote cp` a progress bar and percentage counter are temporarily shown.
---
 tools/mpremote/mpremote/main.py | 29 +++++++++++++++++++++++++++--
 tools/pyboard.py                | 20 ++++++++++++++++----
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/tools/mpremote/mpremote/main.py b/tools/mpremote/mpremote/main.py
index 524b2ec80a..50ad73cc17 100644
--- a/tools/mpremote/mpremote/main.py
+++ b/tools/mpremote/mpremote/main.py
@@ -268,6 +268,27 @@ def do_disconnect(pyb):
     pyb.close()
 
 
+def show_progress_bar(size, total_size):
+    if not sys.stdout.isatty():
+        return
+    verbose_size = 2048
+    bar_length = 20
+    if total_size < verbose_size:
+        return
+    elif size >= total_size:
+        # Clear progress bar when copy completes
+        print("\r" + " " * (20 + bar_length) + "\r", end="")
+    else:
+        progress = size / total_size
+        bar = round(progress * bar_length)
+        print(
+            "\r ... copying {:3.0f}% [{}{}]".format(
+                progress * 100, "#" * bar, "-" * (bar_length - bar)
+            ),
+            end="",
+        )
+
+
 def do_filesystem(pyb, args):
     def _list_recursive(files, path):
         if os.path.isdir(path):
@@ -293,9 +314,13 @@ def do_filesystem(pyb, args):
                 if d not in known_dirs:
                     pyb.exec_("try:\n uos.mkdir('%s')\nexcept OSError as e:\n print(e)" % d)
                     known_dirs.add(d)
-            pyboard.filesystem_command(pyb, ["cp", "/".join((dir, file)), ":" + dir + "/"])
+            pyboard.filesystem_command(
+                pyb,
+                ["cp", "/".join((dir, file)), ":" + dir + "/"],
+                progress_callback=show_progress_bar,
+            )
     else:
-        pyboard.filesystem_command(pyb, args)
+        pyboard.filesystem_command(pyb, args, progress_callback=show_progress_bar)
     args.clear()
 
 
diff --git a/tools/pyboard.py b/tools/pyboard.py
index 15ddfd745d..a9a4cbff87 100755
--- a/tools/pyboard.py
+++ b/tools/pyboard.py
@@ -491,7 +491,10 @@ class Pyboard:
         )
         self.exec_(cmd, data_consumer=stdout_write_bytes)
 
-    def fs_get(self, src, dest, chunk_size=256):
+    def fs_get(self, src, dest, chunk_size=256, progress_callback=None):
+        if progress_callback:
+            src_size = int(self.exec_("import os\nprint(os.stat('%s')[6])" % src))
+            written = 0
         self.exec_("f=open('%s','rb')\nr=f.read" % src)
         with open(dest, "wb") as f:
             while True:
@@ -507,9 +510,15 @@ class Pyboard:
                 if not data:
                     break
                 f.write(data)
+                if progress_callback:
+                    written += len(data)
+                    progress_callback(written, src_size)
         self.exec_("f.close()")
 
-    def fs_put(self, src, dest, chunk_size=256):
+    def fs_put(self, src, dest, chunk_size=256, progress_callback=None):
+        if progress_callback:
+            src_size = os.path.getsize(src)
+            written = 0
         self.exec_("f=open('%s','wb')\nw=f.write" % dest)
         with open(src, "rb") as f:
             while True:
@@ -520,6 +529,9 @@ class Pyboard:
                     self.exec_("w(b" + repr(data) + ")")
                 else:
                     self.exec_("w(" + repr(data) + ")")
+                if progress_callback:
+                    written += len(data)
+                    progress_callback(written, src_size)
         self.exec_("f.close()")
 
     def fs_mkdir(self, dir):
@@ -546,7 +558,7 @@ def execfile(filename, device="/dev/ttyACM0", baudrate=115200, user="micro", pas
     pyb.close()
 
 
-def filesystem_command(pyb, args):
+def filesystem_command(pyb, args, progress_callback=None):
     def fname_remote(src):
         if src.startswith(":"):
             src = src[1:]
@@ -579,7 +591,7 @@ def filesystem_command(pyb, args):
                 src = fname_remote(src)
                 dest2 = fname_cp_dest(src, dest)
                 print(fmt % (src, dest2))
-                op(src, dest2)
+                op(src, dest2, progress_callback=progress_callback)
         else:
             op = {
                 "ls": pyb.fs_ls,