From 41684b9f864e86ed7a6cb07c93469d17e993631c Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Mon, 19 Aug 2024 15:33:11 -0400 Subject: [PATCH 1/3] Implement existing scan commands on top of `callback_for_commands` --- bellows/ezsp/__init__.py | 48 ++++++++++++++++++++++++++++------------ 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py index 273add3b..c199530b 100644 --- a/bellows/ezsp/__init__.py +++ b/bellows/ezsp/__init__.py @@ -224,28 +224,44 @@ async def _list_command( self, name, item_frames, completion_frame, spos, *args, **kwargs ): """Run a command, returning result callbacks as a list""" - fut = asyncio.Future() + queue = asyncio.Queue() results = [] - def cb(frame_name, response): - if frame_name in item_frames: + with self.callback_for_commands( + commands=set(item_frames) | {completion_frame}, + callback=lambda command, response: queue.put_nowait((command, response)), + ): + v = await self._command(name, *args, **kwargs) + if t.sl_Status.from_ember_status(v[0]) != t.sl_Status.OK: + raise Exception(v) + + while True: + command, response = await queue.get() + if command == completion_frame: + if t.sl_Status.from_ember_status(response[spos]) != t.sl_Status.OK: + raise Exception(response) + + break + results.append(response) - elif frame_name == completion_frame: - fut.set_result(response) + + return results + + @contextlib.contextmanager + def callback_for_commands( + self, commands: set[str], callback: Callable + ) -> Generator[None]: + def cb(frame_name, response): + if frame_name in commands: + callback(frame_name, response) cbid = self.add_callback(cb) + try: - v = await self._command(name, *args, **kwargs) - if t.sl_Status.from_ember_status(v[0]) != t.sl_Status.OK: - raise Exception(v) - v = await fut - if t.sl_Status.from_ember_status(v[spos]) != t.sl_Status.OK: - raise Exception(v) + yield finally: self.remove_callback(cbid) - return results - startScan = functools.partialmethod( _list_command, "startScan", @@ -254,7 +270,11 @@ def cb(frame_name, response): 1, ) pollForData = functools.partialmethod( - _list_command, "pollForData", ["pollHandler"], "pollCompleteHandler", 0 + _list_command, + "pollForData", + ["pollHandler"], + "pollCompleteHandler", + 0, ) zllStartScan = functools.partialmethod( _list_command, From d97058f31e170b75c7f45be1b1bd25f0a6f8e521 Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Mon, 19 Aug 2024 15:33:46 -0400 Subject: [PATCH 2/3] Implement `network_scan` --- bellows/zigbee/application.py | 38 +++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 22247178..35834f91 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -5,6 +5,7 @@ import os import statistics import sys +from typing import AsyncGenerator if sys.version_info[:2] < (3, 11): from async_timeout import timeout as asyncio_timeout # pragma: no cover @@ -698,6 +699,43 @@ async def energy_scan( for channel in list(channels) } + async def network_scan( + self, channels: t.Channels, duration: int + ) -> AsyncGenerator[zigpy.types.NetworkBeacon]: + """Scans for networks and yields network beacons.""" + queue = asyncio.Queue() + + with self._ezsp.callback_for_commands( + {"networkFoundHandler", "scanCompleteHandler"}, + callback=lambda command, response: queue.put_nowait((command, response)), + ): + # XXX: replace with normal command invocation once overload is removed + (status,) = await self._ezsp._command( + "startScan", + scanType=t.EzspNetworkScanType.ACTIVE_SCAN, + channelMask=channels, + duration=duration, + ) + + while True: + command, response = await queue.get() + + if command == "scanCompleteHandler": + break + + (networkFound, lastHopLqi, lastHopRssi) = response + + yield zigpy.types.NetworkBeacon( + pan_id=networkFound.panId, + extended_pan_id=networkFound.extendedPanId, + channel=networkFound.channel, + nwk_update_id=networkFound.nwkUpdateId, + permit_joining=bool(networkFound.allowingJoin), + stack_profile=networkFound.stackProfile, + lqi=lastHopLqi, + rssi=lastHopRssi, + ) + async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None: if not self.is_controller_running: raise ControllerError("ApplicationController is not running") From 04550f12b6e4db97394ead834c6e239fae7cb2a7 Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Mon, 19 Aug 2024 15:42:05 -0400 Subject: [PATCH 3/3] `duration` -> `duration_exp` --- bellows/zigbee/application.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 35834f91..554b5985 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -700,7 +700,7 @@ async def energy_scan( } async def network_scan( - self, channels: t.Channels, duration: int + self, channels: t.Channels, duration_exp: int ) -> AsyncGenerator[zigpy.types.NetworkBeacon]: """Scans for networks and yields network beacons.""" queue = asyncio.Queue() @@ -714,7 +714,7 @@ async def network_scan( "startScan", scanType=t.EzspNetworkScanType.ACTIVE_SCAN, channelMask=channels, - duration=duration, + duration=duration_exp, ) while True: