Implement network scanning using the standard zigpy interface

bellows/ezsp/__init__.py

@@ -224,28 +224,44 @@
         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 @@
         1,
     )
     pollForData = functools.partialmethod(
-        _list_command, "pollForData", ["pollHandler"], "pollCompleteHandler", 0
+        _list_command,
+        "pollForData",
+        ["pollHandler"],
+        "pollCompleteHandler",
+        0,
     )
     zllStartScan = functools.partialmethod(
         _list_command,

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 @@
             for channel in list(channels)
         }
 
+    async def network_scan(
+        self, channels: t.Channels, duration_exp: 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_exp,
+            )
+
+            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")