feat: simplify examples (#99)

Author: bdraco

docs/usage.md

@@ -11,21 +11,12 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from collections.abc import Callable
-from typing import TypedDict
 
 import habluetooth
-from aioesphomeapi import APIClient, ReconnectLogic
 
-import bleak_esphome
-
-
-class ESPHomeDeviceConfig(TypedDict):
-    """Configuration for an ESPHome device."""
-
-    address: str
-    noise_psk: str | None
+from bleak_esphome import APIConnectionManager, ESPHomeDeviceConfig
 
+CONNECTION_TIMEOUT = 5
 
 # An unlimited number of devices can be added here
 ESPHOME_DEVICES: list[ESPHomeDeviceConfig] = [
@@ -40,65 +31,36 @@ ESPHOME_DEVICES: list[ESPHomeDeviceConfig] = [
 ]
 
 
-async def setup_api_connection(
-    address: str, noise_psk: str | None = None
-) -> tuple[ReconnectLogic, APIClient]:
-    """Setup the API connection."""
-    cli = APIClient(address=address, port=6053, password=None, noise_psk=noise_psk)
-    unregister_scanner: Callable[[], None] | None = None
-
-    async def on_disconnect(expected_disconnect: bool) -> None:
-        nonlocal unregister_scanner
-        if unregister_scanner is not None:
-            unregister_scanner()
-            unregister_scanner = None
-
-    async def on_connect() -> None:
-        nonlocal unregister_scanner
-        device_info = await cli.device_info()
-        client_data = bleak_esphome.connect_scanner(cli, device_info, True)
-        scanner = client_data.scanner
-        assert scanner is not None  # noqa: S101
-        scanner.async_setup()
-        unregister_scanner = habluetooth.get_manager().async_register_scanner(scanner)
-
-    reconnect_logic = ReconnectLogic(
-        client=cli,
-        on_disconnect=on_disconnect,
-        on_connect=on_connect,
-    )
-    await reconnect_logic.start()
-
-    return reconnect_logic, cli
-
+async def example_app() -> None:
+    """Example application here."""
+    import bleak
 
-async def run_application(cli: APIClient) -> None:
-    """Test application here."""
-    import bleak  # noqa
+    await asyncio.sleep(5) # Give time for advertisements to be received
 
     # Use bleak normally here
+    devices = await bleak.BleakScanner.discover(return_adv=True)
+    for d, a in devices.values():
+        print()
+        print(d)
+        print("-" * len(str(d)))
+        print(a)
 
     # Wait forever
-    event = asyncio.Event()
-    await event.wait()
+    await asyncio.Event().wait()
 
 
 async def run() -> None:
     """Run the main application."""
-    esphome_connections: list[tuple[ReconnectLogic, APIClient]] = []
-    reconnect_logic: ReconnectLogic | None = None
-    cli: APIClient | None = None
+    connections = [APIConnectionManager(device) for device in ESPHOME_DEVICES]
+    await habluetooth.BluetoothManager().async_setup()
     try:
-        await habluetooth.BluetoothManager().async_setup()
-        for device in ESPHOME_DEVICES:
-            esphome_connections.append(
-                await setup_api_connection(device["address"], device["noise_psk"])
-            )
-        await run_application(cli)
+        await asyncio.wait(
+            (asyncio.create_task(conn.start()) for conn in connections),
+            timeout=CONNECTION_TIMEOUT,
+        )
+        await example_app()
     finally:
-        for reconnect_logic, cli in esphome_connections:
-            await reconnect_logic.stop()
-            await cli.disconnect()
+        await asyncio.gather(*(conn.stop() for conn in connections))
 
 
 logging.basicConfig(level=logging.DEBUG)

examples/setup_scanner.py

@@ -2,21 +2,12 @@
 
 import asyncio
 import logging
-from collections.abc import Callable
-from typing import TypedDict
 
 import habluetooth
-from aioesphomeapi import APIClient, ReconnectLogic
 
-import bleak_esphome
-
-
-class ESPHomeDeviceConfig(TypedDict):
-    """Configuration for an ESPHome device."""
-
-    address: str
-    noise_psk: str | None
+from bleak_esphome import APIConnectionManager, ESPHomeDeviceConfig
 
+CONNECTION_TIMEOUT = 5
 
 # An unlimited number of devices can be added here
 ESPHOME_DEVICES: list[ESPHomeDeviceConfig] = [
@@ -31,65 +22,36 @@ class ESPHomeDeviceConfig(TypedDict):
 ]
 
 
-async def setup_api_connection(
-    address: str, noise_psk: str | None = None
-) -> tuple[ReconnectLogic, APIClient]:
-    """Setup the API connection."""
-    cli = APIClient(address=address, port=6053, password=None, noise_psk=noise_psk)
-    unregister_scanner: Callable[[], None] | None = None
-
-    async def on_disconnect(expected_disconnect: bool) -> None:
-        nonlocal unregister_scanner
-        if unregister_scanner is not None:
-            unregister_scanner()
-            unregister_scanner = None
-
-    async def on_connect() -> None:
-        nonlocal unregister_scanner
-        device_info = await cli.device_info()
-        client_data = bleak_esphome.connect_scanner(cli, device_info, True)
-        scanner = client_data.scanner
-        assert scanner is not None  # noqa: S101
-        scanner.async_setup()
-        unregister_scanner = habluetooth.get_manager().async_register_scanner(scanner)
-
-    reconnect_logic = ReconnectLogic(
-        client=cli,
-        on_disconnect=on_disconnect,
-        on_connect=on_connect,
-    )
-    await reconnect_logic.start()
-
-    return reconnect_logic, cli
-
+async def example_app() -> None:
+    """Example application here."""
+    import bleak
 
-async def run_application(cli: APIClient) -> None:
-    """Test application here."""
-    import bleak  # noqa
+    await asyncio.sleep(5)  # Give time for the scanner to find devices
 
     # Use bleak normally here
+    devices = await bleak.BleakScanner.discover(return_adv=True)
+    for d, a in devices.values():
+        print()
+        print(d)
+        print("-" * len(str(d)))
+        print(a)
 
     # Wait forever
-    event = asyncio.Event()
-    await event.wait()
+    await asyncio.Event().wait()
 
 
 async def run() -> None:
     """Run the main application."""
-    esphome_connections: list[tuple[ReconnectLogic, APIClient]] = []
-    reconnect_logic: ReconnectLogic | None = None
-    cli: APIClient | None = None
+    connections = [APIConnectionManager(device) for device in ESPHOME_DEVICES]
+    await habluetooth.BluetoothManager().async_setup()
     try:
-        await habluetooth.BluetoothManager().async_setup()
-        for device in ESPHOME_DEVICES:
-            esphome_connections.append(
-                await setup_api_connection(device["address"], device["noise_psk"])
-            )
-        await run_application(cli)
+        await asyncio.wait(
+            (asyncio.create_task(conn.start()) for conn in connections),
+            timeout=CONNECTION_TIMEOUT,
+        )
+        await example_app()
     finally:
-        for reconnect_logic, cli in esphome_connections:
-            await reconnect_logic.stop()
-            await cli.disconnect()
+        await asyncio.gather(*(conn.stop() for conn in connections))
 
 
 logging.basicConfig(level=logging.DEBUG)

poetry.lock

@@ -31,7 +31,7 @@ python = ">=3.11,<3.14"
 aioesphomeapi = ">=29"
 bleak = ">=0.21.1"
 bluetooth-data-tools = ">=1.18.0"
-habluetooth = ">=3.17.0"
+habluetooth = ">=3.24.1"
 lru-dict = ">=1.2.0"
 bleak-retry-connector = ">=3.8.0"
 

pyproject.toml

@@ -1,5 +1,8 @@
 from .connect import connect_scanner
+from .connection_manager import APIConnectionManager, ESPHomeDeviceConfig
 
 __all__ = [
+    "APIConnectionManager",
+    "ESPHomeDeviceConfig",
     "connect_scanner",
 ]

src/bleak_esphome/__init__.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from typing import TypedDict
+
+import habluetooth
+from aioesphomeapi import APIClient, ReconnectLogic
+
+import bleak_esphome
+
+
+class ESPHomeDeviceConfig(TypedDict):
+    """Configuration for an ESPHome device."""
+
+    address: str
+    noise_psk: str | None
+
+
+class APIConnectionManager:
+    """Manager for the API connection to an ESPHome device."""
+
+    def __init__(self, config: ESPHomeDeviceConfig) -> None:
+        """Initialize the API connection manager."""
+        self._address = config["address"]
+        self._noise_psk = config["noise_psk"]
+        self._cli: APIClient = APIClient(
+            address=self._address, port=6053, password=None, noise_psk=self._noise_psk
+        )
+        self._reconnect_logic = ReconnectLogic(
+            client=self._cli,
+            on_disconnect=self._on_disconnect,
+            on_connect=self._on_connect,
+        )
+        self._unregister_scanner: Callable[[], None] | None = None
+        self._start_future: asyncio.Future[None] = (
+            asyncio.get_running_loop().create_future()
+        )
+
+    async def _on_disconnect(self, expected_disconnect: bool) -> None:
+        """Handle the disconnection of the API client."""
+        if self._unregister_scanner is not None:
+            self._unregister_scanner()
+            self._unregister_scanner = None
+
+    async def _on_connect(self) -> None:
+        """Handle the connection of the API client."""
+        device_info = await self._cli.device_info()
+        client_data = bleak_esphome.connect_scanner(self._cli, device_info, True)
+        scanner = client_data.scanner
+        assert scanner is not None  # noqa: S101
+        scanner.async_setup()
+        self._unregister_scanner = habluetooth.get_manager().async_register_scanner(
+            scanner
+        )
+        if not self._start_future.done():
+            self._start_future.set_result(None)
+
+    async def start(self) -> None:
+        """Start the API connection."""
+        await self._reconnect_logic.start()
+        await self._start_future
+
+    async def stop(self) -> None:
+        """Stop the API connection."""
+        await self._reconnect_logic.stop()
+        await self._cli.disconnect()
+        if not self._start_future.done():
+            self._start_future.cancel()
+        if self._unregister_scanner is not None:
+            self._unregister_scanner()
+            self._unregister_scanner = None