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")