Do this without the asyncio with.

ch-13/1-modelling-space/computer/display_arena_from_robot.py

@@ -9,7 +9,7 @@ from robot_ble_connection import BleConnection
 class RobotDisplay:
     def __init__(self):
         self.fig, self.ax = plt.subplots()
-        self.ble_connection = BleConnection(self.handle_data, self.connected)
+        self.ble_connection = BleConnection(self.handle_data)
 
         self.arena = None
         self.display_closed = False
@@ -35,29 +35,32 @@ class RobotDisplay:
         self.arena = arena
         self.ax.clear()
         for line in arena["arena"]:
-            self.ax.plot([line[0][0], line[1][0]],
-                         [line[0][1], line[1][1]],
-                         color='black')
+            self.ax.plot(
+                [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="black"
+            )
         for line in arena["target_zone"]:
-            self.ax.plot([line[0][0], line[1][0]],
-                         [line[0][1], line[1][1]],
-                         color="red")
-
-    def connected(self):
-        request = json.dumps({"command": "arena"}).encode()
-        print(f"Sending request for arena: {request}")
-        self.ble_connection.send_uart_data(request)
+            self.ax.plot(
+                [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="red"
+            )
 
     async def main(self):
         plt.ion()
 
-        asyncio.create_task(self.ble_connection.connect())
+        try:
+            await self.ble_connection.connect()
+            request = json.dumps({"type": "arena"}).encode()
+            print(f"Sending request for arena: {request}")
+            self.ble_connection.send_uart_data(request)
+
             while not self.display_closed:
                 plt.pause(0.05)
                 plt.draw()
                 await asyncio.sleep(0.01)
 
             plt.show()
+        finally:
+            await self.ble_connection.close()
+
 
 robot_display = RobotDisplay()
 asyncio.run(robot_display.main())

ch-13/1-modelling-space/computer/robot_ble_connection.py

@@ -11,11 +11,9 @@ class BleConnection:
     adafruit_tx_uuid = "6E400002-B5A3-F393-E0A9-E50E24DCCA9E"
     ble_name = "Adafruit Bluefruit LE"
 
-    def __init__(self, receive_handler: typing.Callable[[bytes], None], 
-            connected_handler: typing.Callable[[], None]):
+    def __init__(self, receive_handler: typing.Callable[[bytes], None]):
         self.ble_client : bleak.BleakClient = None
         self.receive_handler = receive_handler
-        self.connected_handler = connected_handler
 
     def _uart_handler(self, _, data: bytes):
         self.receive_handler(data)
@@ -27,15 +25,15 @@ class BleConnection:
         print([device.name for device in devices])        
         ble_device_info = [device for device in devices if device.name==self.ble_name][0]
         print("Connecting to {}...".format(ble_device_info.name))
-        async with bleak.BleakClient(ble_device_info.address) as ble_client:
-            self.ble_client = ble_client
-            self.connected_handler()
+        self.ble_client = bleak.BleakClient(ble_device_info.address)
+        await self.ble_client.connect()
         print("Connected to {}".format(ble_device_info.name))
-            asyncio.create_task(
+        self.notify_task = asyncio.create_task(
             self.ble_client.start_notify(self.adafruit_rx_uuid, self._uart_handler)
         )
-            while True:
-                await asyncio.sleep(1)
+
+    async def close(self):
+        await self.ble_client.disconnect()
                 
     def send_uart_data(self, data: bytes):
         if self.ble_client: