tools/mpr: Improve device auto connection.

dpgeorge · dpgeorge · commit a6d01fea5bfb · 2021-05-06T15:38:51.000+10:00
Signed-off-by: Damien George &lt;damien@micropython.org&gt;
diff --git a/tools/mpr b/tools/mpr
@@ -583,8 +583,9 @@ class SerialIntercept:
 
 
 class PyboardExtended(pyboard.Pyboard):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, dev, *args, **kwargs):
+        super().__init__(dev, *args, **kwargs)
+        self.device_name = dev
         self.mounted = False
 
     def enter_raw_repl_without_soft_reset(self):
@@ -652,22 +653,6 @@ class PyboardExtended(pyboard.Pyboard):
             self.mounted = False
 
 
-def find_serial_device(dev):
-    ports = serial.tools.list_ports.comports()
-    if dev is None:
-        for dev in device_shortcuts.values():
-            if any(p.device == dev for p in ports):
-                return dev
-        print("no device found")
-        sys.exit(1)
-    else:
-        dev = device_shortcuts.get(dev, dev)
-        if any(p.device == dev for p in ports):
-            return dev
-        print(f"{dev} could not be accessed")
-        sys.exit(1)
-
-
 class ConsolePosix:
     def __init__(self):
         self.infd = sys.stdin.fileno()
@@ -814,6 +799,33 @@ else:
         VT_ENABLED = False
 
 
+def do_connect(args):
+    try:
+        if args and (args[0] == "connect" or args[0] in device_shortcuts):
+            if args[0] == "connect":
+                args.pop(0)
+            dev = args.pop(0)
+            dev = device_shortcuts.get(dev, dev)
+            return PyboardExtended(dev, baudrate=115200)
+        else:
+            # auto-detect and auto-connect
+            ports = serial.tools.list_ports.comports()
+            for dev in device_shortcuts.values():
+                if any(p.device == dev for p in ports):
+                    try:
+                        return PyboardExtended(dev, baudrate=115200)
+                    except pyboard.PyboardError as er:
+                        if not er.args[0].startswith("failed to access"):
+                            raise er
+            raise pyboard.PyboardError("no device found")
+    except pyboard.PyboardError as er:
+        msg = er.args[0]
+        if msg.startswith("failed to access"):
+            msg += " (it may be in use by another program)"
+        print(msg)
+        sys.exit(1)
+
+
 def do_filesystem(pyb, args):
     def _list_recursive(files, path):
         if os.path.isdir(path):
@@ -892,7 +904,7 @@ def do_repl_main_loop(pyb, console_in, console_out_write, file_to_inject):
                     console_out_write(b"[%02x]" % ord(c))
 
 
-def do_repl(pyb, dev, args):
+def do_repl(pyb, args):
     if len(args) and args[0] == "--capture":
         args.pop(0)
         capture_file = args.pop(0)
@@ -901,7 +913,7 @@ def do_repl(pyb, dev, args):
 
     file_to_inject = args.pop(0) if len(args) else None
 
-    print("Connected to MicroPython at %s" % dev)
+    print("Connected to MicroPython at %s" % pyb.device_name)
     print("Use Ctrl-] to exit this shell")
     if capture_file is not None:
         print('Capturing session to file "%s"' % capture_file)
@@ -946,24 +958,7 @@ def execbuffer(pyb, buf, follow):
 def main():
     args = sys.argv[1:]
 
-    if args and args[0] in device_shortcuts:
-        dev = find_serial_device(args.pop(0))
-    elif args and args[0] == "connect":
-        args.pop(0)
-        dev = find_serial_device(args.pop(0))
-    else:
-        # auto-detect and auto-connect
-        dev = find_serial_device(None)
-
-    try:
-        pyb = PyboardExtended(dev, baudrate=115200)
-    except pyboard.PyboardError as er:
-        msg = er.args[0]
-        if msg.startswith("failed to access"):
-            msg += " (it may be in use by another program)"
-        print(msg)
-        sys.exit(1)
-
+    pyb = do_connect(args)
     in_raw_repl = False
     did_action = False
 
@@ -1018,7 +1013,7 @@ def main():
             elif cmd == "fs":
                 do_filesystem(pyb, args)
             elif cmd == "repl":
-                do_repl(pyb, dev, args)
+                do_repl(pyb, args)
             elif cmd == "run":
                 follow = True
                 if args[0] == "--no-follow":
@@ -1039,7 +1034,7 @@ def main():
             if in_raw_repl:
                 pyb.exit_raw_repl()
                 in_raw_repl = False
-            do_repl(pyb, dev, args)
+            do_repl(pyb, args)
     finally:
         if pyb is not None:
             if pyb.mounted:
