Slim device list

mdeweerd · Apr 28, 2023 · 4de930a · 4de930a
1 parent dfff5c2
commit 4de930a
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 50 deletions.
diff --git a/README.md b/README.md
@@ -113,7 +113,7 @@ ZHA Toolkit can also:
     - [`backup`: Backup the coordinator](#backup-backup-the-coordinator)
     - [`misc_settime`: Set attributes of a Time Cluster](#misc_settime-set-attributes-of-a-time-cluster)
     - [`ota_notify` - Download/Trigger Device FW update](#ota_notify---downloadtrigger-device-fw-update)
-    - [`zha_devices`: Device List Information to Event or CSV](#zha_devices-device-list-information-to-event-or-csv)
+    - [`zha_devices`: Device Information to Event or CSV](#zha_devices-device-information-to-event-or-csv)
     - [`register_services`: Reregister ZHA-Toolkit services](#register_services-reregister-zha-toolkit-services)
     - [`ha_set_state` - Update HA state](#ha_set_state---update-ha-state)
   - [User method](#user-method)
@@ -1687,11 +1687,17 @@ data:
   path: /config/zb_ota
 ```
 
-### `zha_devices`: Device List Information to Event or CSV
+### `zha_devices`: Device Information to Event or CSV
 
-Write information from currently known ZHA devices to a CSV file. You also
-get this data in the 'devices' field of the generated events which allows
-you to get information about endpoints and services as well.
+Get Device information as event data or in a CSV file.
+
+You can select the data fields in the CSV and the event data through the
+command_data parameter. If you do not provide a list, a default list is
+used for the CSV file, and all available data is provided in the devices
+field of the event data.
+
+You also get this data in the 'devices' field of the generated events which
+allows you to get information about endpoints and services as well.
 
 ```yaml
 service: zha_toolkit.zha_devices

diff --git a/custom_components/zha_toolkit/zha.py b/custom_components/zha_toolkit/zha.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Any
 
 from . import utils as u
 from .params import INTERNAL_PARAMS as p
@@ -11,8 +12,43 @@
 async def zha_devices(
     app, listener, ieee, cmd, data, service, params, event_data
 ):
+    doGenerateCSV = params[p.CSV_FILE] is not None
+
+    # Determine fields to render.
+    # If the user provides a list, it is also used to
+    # limit the contents of "devices" in the event_data.
+    if data is not None and isinstance(data, list):
+        selectDeviceFields = True
+        columns = data
+    else:
+        selectDeviceFields = False
+        columns = [
+            "ieee",
+            "nwk",
+            "manufacturer",
+            "model",
+            "name",
+            "quirk_applied",
+            "quirk_class",
+            "manufacturer_code",
+            "power_source",
+            "lqi",
+            "rssi",
+            "last_seen",
+            "available",
+            "device_type",
+            "user_given_name",
+            "device_reg_id",
+            "area_id",
+        ]
+        # TODO: Skipped in columns, needs special handling
+        #  'signature'
+        #  'endpoints'
 
     devices = [device.zha_device_info for device in listener.devices.values()]
+    # Set default value for 'devices' in event_data,
+    # may be slimmed down.  Ensures that devices is set in case
+    # an exception occurs.
     event_data["devices"] = devices
 
     if params[p.CSV_LABEL] is not None and isinstance(
@@ -33,59 +69,47 @@ async def zha_devices(
         except Exception:  # nosec
             pass
 
-    if params[p.CSV_FILE] is not None:
-        if data is not None and isinstance(data, list):
-            columns = data
-        else:
-            columns = [
-                "ieee",
-                "nwk",
-                "manufacturer",
-                "model",
-                "name",
-                "quirk_applied",
-                "quirk_class",
-                "manufacturer_code",
-                "power_source",
-                "lqi",
-                "rssi",
-                "last_seen",
-                "available",
-                "device_type",
-                "user_given_name",
-                "device_reg_id",
-                "area_id",
-            ]
-            # TODO: Skipped in columns, needs special handling
-            #  'signature'
-            #  'endpoints'
+    if doGenerateCSV or selectDeviceFields:
 
-        u.append_to_csvfile(
-            columns,
-            "csv",
-            params[p.CSV_FILE],
-            "device_dump['HEADER']",
-            listener=listener,
-            overwrite=True,
-        )
+        if doGenerateCSV:
+            # Write CSV header
+            u.append_to_csvfile(
+                columns,
+                "csv",
+                params[p.CSV_FILE],
+                "device_dump['HEADER']",
+                listener=listener,
+                overwrite=True,
+            )
 
+        slimmedDevices: list[Any] = []
         for d in devices:
-            fields: list[int | str | None] = []
+            # Fields for CSV
+            csvFields: list[int | str | None] = []
+            # Fields for slimmed devices dict
+            rawFields: dict[str, Any] = {}
+
             for c in columns:
                 if c not in d.keys():
-                    fields.append(None)
+                    csvFields.append(None)
                 else:
                     val = d[c]
+                    rawFields[c] = val
                     if c in ["manufacturer", "nwk"] and isinstance(val, int):
                         val = f"0x{val:04X}"
 
-                    fields.append(d[c])
+                    csvFields.append(d[c])
 
-            LOGGER.debug("Device %r", fields)
-            u.append_to_csvfile(
-                fields,
-                "csv",
-                params[p.CSV_FILE],
-                f"device_dump[{d['ieee']}]",
-                listener=listener,
-            )
+            slimmedDevices.append(rawFields)
+
+            if doGenerateCSV:
+                LOGGER.debug("Device %r", csvFields)
+                u.append_to_csvfile(
+                    csvFields,
+                    "csv",
+                    params[p.CSV_FILE],
+                    f"device_dump[{d['ieee']}]",
+                    listener=listener,
+                )
+        if selectDeviceFields:
+            event_data["devices"] = slimmedDevices