pybricksdev.connections.pybricks.PybricksHub: fix connection state management

PybricksHub.disconnect_observable was one-use only since it used an
AsyncSubject. This caused the PybricksHub.race_disconnect method to
think that the hub was still disconnected after a reconnection.

To fix this, we replace PybricksHub.disconnect_observable and
PybricksHub.connected with a new PybricksHub.connection_state_observable
that uses a BehaviorSubject so that it can be updated across multiple
connect/disconnect events.

This also introduces a new pybricksdev.connections.ConnectionState enum
that give a more nuanced state indication rather than using a bool.

Fixes: pybricks/support#971

Author: dlech

CHANGELOG.md

@@ -6,6 +6,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Added `pybricksdev.connections.ConnectionState` enum class.
+- Added `pybricksdev.connections.pybricks.PybricksHub.connection_state_observable` attribute.
+
+### Fixed
+- Fixed `pybricksdev.connections.pybricks.PybricksHub` disconnect state not reset after reconnect ([support#971]).
+
+### Removed
+- Removed `pybricksdev.connections.pybricks.PybricksHub.disconnect_observable` attribute.
+- Removed `pybricksdev.connections.pybricks.PybricksHub.connected` attribute.
+
 ## [1.0.0-alpha.37] - 2023-02-27
 
 ### Added

pybricksdev/connections/__init__.py

@@ -1,2 +1,27 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2022 The Pybricks Authors
+# Copyright (c) 2023 The Pybricks Authors
+
+import enum
+
+
+class ConnectionState(enum.Enum):
+    """
+    Indicates state of a connection.
+    """
+
+    CONNECTING = enum.auto()
+    """
+    The device is in the process of connecting.
+    """
+    CONNECTED = enum.auto()
+    """
+    The device is connected.
+    """
+    DISCONNECTING = enum.auto()
+    """
+    The device is in the process of disconnecting.
+    """
+    DISCONNECTED = enum.auto()
+    """
+    The device is disconnected.
+    """

pybricksdev/connections/pybricks.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2021-2022 The Pybricks Authors
+# Copyright (c) 2021-2023 The Pybricks Authors
 
 import asyncio
+import contextlib
 import logging
 import os
 import struct
@@ -12,7 +13,7 @@
 from bleak import BleakClient
 from bleak.backends.device import BLEDevice
 from packaging.version import Version
-from rx.subject import AsyncSubject, BehaviorSubject, Subject
+from rx.subject import BehaviorSubject, Subject
 from tqdm.auto import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
 
@@ -35,6 +36,7 @@
 from ..compile import compile_file, compile_multi_file
 from ..tools import chunk
 from ..tools.checksum import xor_bytes
+from . import ConnectionState
 
 logger = logging.getLogger(__name__)
 
@@ -76,7 +78,7 @@ class PybricksHub:
     """
 
     def __init__(self):
-        self.disconnect_observable = AsyncSubject()
+        self.connection_state_observable = BehaviorSubject(ConnectionState.DISCONNECTED)
         self.status_observable = BehaviorSubject(StatusFlag(0))
         self.nus_observable = Subject()
         self.stream_buf = bytearray()
@@ -87,9 +89,6 @@ def __init__(self):
         self._capability_flags = HubCapabilityFlag(0)
         self._max_user_program_size = 0
 
-        # indicates that the hub is currently connected via BLE
-        self.connected = False
-
         # indicates is we are currently downloading a program over NUS (legacy download)
         self._downloading_via_nus = False
 
@@ -188,17 +187,28 @@ async def connect(self, device: BLEDevice):
         """
         logger.info(f"Connecting to {device.name}")
 
-        def handle_disconnect(client: BleakClient):
-            logger.info("Disconnected!")
-            self.disconnect_observable.on_next(True)
-            self.disconnect_observable.on_completed()
-            self.connected = False
+        if self.connection_state_observable.value != ConnectionState.DISCONNECTED:
+            raise RuntimeError(
+                f"attempting to connect with invalid state: {self.connection_state_observable.value}"
+            )
 
-        self.client = BleakClient(device, disconnected_callback=handle_disconnect)
+        async with contextlib.AsyncExitStack() as stack:
+            self.connection_state_observable.on_next(ConnectionState.CONNECTING)
 
-        await self.client.connect()
+            stack.callback(
+                self.connection_state_observable.on_next, ConnectionState.DISCONNECTED
+            )
+
+            def handle_disconnect(_: BleakClient):
+                logger.info("Disconnected!")
+                self.connection_state_observable.on_next(ConnectionState.DISCONNECTED)
+
+            self.client = BleakClient(device, disconnected_callback=handle_disconnect)
+
+            await self.client.connect()
+
+            stack.push_async_callback(self.disconnect)
 
-        try:
             logger.info("Connected successfully!")
 
             fw_version = await self.client.read_gatt_char(FW_REV_UUID)
@@ -236,17 +246,24 @@ def handle_disconnect(client: BleakClient):
             await self.client.start_notify(
                 PYBRICKS_COMMAND_EVENT_UUID, self.pybricks_service_handler
             )
-            self.connected = True
-        except:  # noqa: E722
-            self.disconnect()
-            raise
+
+            self.connection_state_observable.on_next(ConnectionState.CONNECTED)
+
+            # don't unwind on success
+            stack.pop_all()
 
     async def disconnect(self):
-        if self.connected:
-            logger.info("Disconnecting...")
+        logger.info("Disconnecting...")
+
+        if self.connection_state_observable.value == ConnectionState.CONNECTED:
+            self.connection_state_observable.on_next(ConnectionState.DISCONNECTING)
             await self.client.disconnect()
+            # ConnectionState.DISCONNECTED should be set by disconnect callback
+            assert (
+                self.connection_state_observable.value == ConnectionState.DISCONNECTED
+            )
         else:
-            logger.debug("already disconnected")
+            logger.debug("skipping disconnect because not connected")
 
     async def race_disconnect(self, awaitable: Awaitable[T]) -> T:
         """
@@ -273,7 +290,11 @@ async def race_disconnect(self, awaitable: Awaitable[T]) -> T:
         disconnect_event = asyncio.Event()
         disconnect_task = asyncio.ensure_future(disconnect_event.wait())
 
-        with self.disconnect_observable.subscribe(lambda _: disconnect_event.set()):
+        def handle_disconnect(state: ConnectionState):
+            if state == ConnectionState.DISCONNECTED:
+                disconnect_event.set()
+
+        with self.connection_state_observable.subscribe(handle_disconnect):
             done, pending = await asyncio.wait(
                 {awaitable_task, disconnect_task},
                 return_when=asyncio.FIRST_COMPLETED,
@@ -301,7 +322,7 @@ async def run(
             wait: If true, wait for the user program to stop before returning.
             print_output: If true, echo stdout of the hub to ``sys.stdout``.
         """
-        if not self.connected:
+        if self.connection_state_observable.value != ConnectionState.CONNECTED:
             raise RuntimeError("not connected")
 
         # Reset output buffer