cli.lwp3.repl: end program on disconnect

If the device disconnects (e.g. it is powered off), then the prompt
is still waiting and will cause a crash on the next input.

To work around this, we cancel the task from the disconnect callback
and then catch the exception.

Author: dlech

CHANGELOG.md

@@ -95,6 +95,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Wait for some time to allow program output to be received before disconnecting
   in the `run` command.
 - Fixed spelling of `INPUT` in `pybricksdev.ble.lwp3.messages`.
+- Fixed `pybricksdev lwp3 repl` does not exit if device disconnects.
 
 ## [1.0.0-alpha.3] - 2021-04-09
 ### Changed

pybricksdev/cli/lwp3/repl.py

@@ -104,50 +104,62 @@ def match_lwp3_uuid(dev: BLEDevice, adv: AdvertisementData) -> None:
 
     logger.info("found device")
 
-    def handle_disconnect(client: BleakClient):
-        logger.info("disconnected")
-
-    async with BleakClient(device, disconnected_callback=handle_disconnect) as client:
-        logger.info("connected")
-
-        def handle_notify(handle, value):
-            try:
-                msg = parse_message(value)
-            except Exception as ex:
-                logger.warning("failed to parse message: %s", value, exc_info=ex)
-            else:
-                logger.info("received: %s", msg)
+    repl_task = asyncio.current_task()
 
-        await client.start_notify(LWP3_HUB_CHARACTERISTIC_UUID, handle_notify)
-
-        # welcome is delayed to allow initial log messages to settle.
-        async def welcome():
-            await asyncio.sleep(1)
-            print("Type message and press ENTER to send. Press CTRL+D to exit.")
+    def handle_disconnect(client: BleakClient):
+        repl_task.cancel()
 
-        asyncio.ensure_future(welcome())
+    try:
+        async with BleakClient(
+            device, disconnected_callback=handle_disconnect
+        ) as client:
+            logger.info("connected")
 
-        while True:
-            with patch_stdout():
+            def handle_notify(handle, value):
+                try:
+                    msg = parse_message(value)
+                except Exception as ex:
+                    logger.warning("failed to parse message: %s", value, exc_info=ex)
+                else:
+                    logger.info("received: %s", msg)
+
+            await client.start_notify(LWP3_HUB_CHARACTERISTIC_UUID, handle_notify)
+
+            # welcome is delayed to allow initial log messages to settle.
+            async def welcome():
+                await asyncio.sleep(1)
+                print("Type message and press ENTER to send. Press CTRL+D to exit.")
+
+            asyncio.ensure_future(welcome())
+
+            while True:
+                with patch_stdout():
+                    try:
+                        result = await session.prompt_async(">>> ")
+                    except KeyboardInterrupt:
+                        # CTRL+C ignores the line
+                        continue
+                    except EOFError:
+                        # CTRL+D exits the program
+                        break
                 try:
-                    result = await session.prompt_async(">>> ")
-                except KeyboardInterrupt:
-                    # CTRL+C ignores the line
-                    continue
-                except EOFError:
-                    # CTRL+D exits the program
-                    break
-            try:
-                msg = eval(result, _eval_pool)
-                if not isinstance(msg, AbstractMessage):
-                    raise ValueError("not a message object")
-            except Exception:
-                logger.exception("bad input:")
-            else:
-                logger.info("sending: %s", msg)
-                await client.write_gatt_char(LWP3_HUB_CHARACTERISTIC_UUID, bytes(msg))
-
-        logger.info("disconnecting...")
+                    msg = eval(result, _eval_pool)
+                    if not isinstance(msg, AbstractMessage):
+                        raise ValueError("not a message object")
+                except Exception:
+                    logger.exception("bad input:")
+                else:
+                    logger.info("sending: %s", msg)
+                    await client.write_gatt_char(
+                        LWP3_HUB_CHARACTERISTIC_UUID, bytes(msg)
+                    )
+
+            logger.info("disconnecting...")
+    except asyncio.CancelledError:
+        # happens on disconnect
+        pass
+
+    logger.info("disconnected")
 
 
 def setup_repl_logging() -> None: