1 files changed, 56 insertions, 6 deletions

diff --git a/tools/mpremote/mpremote/main.py b/tools/mpremote/mpremote/main.py
index b30a1a2135..0441857fab 100644
--- a/tools/mpremote/mpremote/main.py
+++ b/tools/mpremote/mpremote/main.py
@@ -181,7 +181,11 @@ def argparse_rtc():
 
 
 def argparse_filesystem():
-    cmd_parser = argparse.ArgumentParser(description="execute filesystem commands on the device")
+    cmd_parser = argparse.ArgumentParser(
+        description="execute filesystem commands on the device",
+        add_help=False,
+    )
+    cmd_parser.add_argument("--help", action="help", help="show this help message and exit")
     _bool_flag(cmd_parser, "recursive", "r", False, "recursive (for cp and rm commands)")
     _bool_flag(
         cmd_parser,
@@ -197,10 +201,26 @@ def argparse_filesystem():
         None,
         "enable verbose output (defaults to True for all commands except cat)",
     )
+    size_group = cmd_parser.add_mutually_exclusive_group()
+    size_group.add_argument(
+        "--size",
+        "-s",
+        default=False,
+        action="store_true",
+        help="show file size in bytes(tree command only)",
+    )
+    size_group.add_argument(
+        "--human",
+        "-h",
+        default=False,
+        action="store_true",
+        help="show file size in a more human readable way (tree command only)",
+    )
+
     cmd_parser.add_argument(
         "command",
         nargs=1,
-        help="filesystem command (e.g. cat, cp, sha256sum, ls, rm, rmdir, touch)",
+        help="filesystem command (e.g. cat, cp, sha256sum, ls, rm, rmdir, touch, tree)",
     )
     cmd_parser.add_argument("path", nargs="+", help="local and remote paths")
     return cmd_parser
@@ -355,10 +375,31 @@ _BUILTIN_COMMAND_EXPANSIONS = {
     "rmdir": "fs rmdir",
     "sha256sum": "fs sha256sum",
     "touch": "fs touch",
+    "tree": "fs tree",
     # Disk used/free.
     "df": [
         "exec",
-        "import os\nprint('mount \\tsize \\tused \\tavail \\tuse%')\nfor _m in [''] + os.listdir('/'):\n _s = os.stat('/' + _m)\n if not _s[0] & 1 << 14: continue\n _s = os.statvfs(_m)\n if _s[0]:\n  _size = _s[0] * _s[2]; _free = _s[0] * _s[3]; print(_m, _size, _size - _free, _free, int(100 * (_size - _free) / _size), sep='\\t')",
+        """
+import os,vfs
+_f = "{:<10}{:>9}{:>9}{:>9}{:>5} {}"
+print(_f.format("filesystem", "size", "used", "avail", "use%", "mounted on"))
+try:
+ _ms = vfs.mount()
+except:
+ _ms = []
+ for _m in [""] + os.listdir("/"):
+  _m = "/" + _m
+  _s = os.stat(_m)
+  if _s[0] & 1 << 14:
+   _ms.append(("<unknown>",_m))
+for _v,_p in _ms:
+ _s = os.statvfs(_p)
+ _sz = _s[0]*_s[2]
+ if _sz:
+  _av = _s[0]*_s[3]
+  _us = 100*(_sz-_av)//_sz
+  print(_f.format(str(_v), _sz, _sz-_av, _av, _us, _p))
+""",
     ],
     # Other shortcuts.
     "reset": {
@@ -552,8 +593,13 @@ def main():
                 command_args = remaining_args
                 extra_args = []
 
-            # Special case: "fs ls" allowed have no path specified.
-            if cmd == "fs" and len(command_args) == 1 and command_args[0] == "ls":
+            # Special case: "fs ls" and "fs tree" can have only options and no path specified.
+            if (
+                cmd == "fs"
+                and len(command_args) >= 1
+                and command_args[0] in ("ls", "tree")
+                and sum(1 for a in command_args if not a.startswith('-')) == 1
+            ):
                 command_args.append("")
 
             # Use the command-specific argument parser.
@@ -574,7 +620,11 @@ def main():
         # If no commands were "actions" then implicitly finish with the REPL
         # using default args.
         if state.run_repl_on_completion():
-            do_repl(state, argparse_repl().parse_args([]))
+            disconnected = do_repl(state, argparse_repl().parse_args([]))
+
+            # Handle disconnection message
+            if disconnected:
+                print("\ndevice disconnected")
 
         return 0
     except CommandError as e: