From e9da9e5df405c5dcbcd9a65b7456a5bb63b23011 Mon Sep 17 00:00:00 2001
From: "Travis F. Collins" <travis.collins@analog.com>
Date: Fri, 1 Nov 2024 07:59:18 -0600
Subject: [PATCH] Allow support for array properties in pybench

Signed-off-by: Travis F. Collins <travis.collins@analog.com>
---
 bench/cli/iiotools.py               | 13 +++++++++++--
 bench/keysight/dwta/data_capture.py | 18 +++++++++++++++---
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/bench/cli/iiotools.py b/bench/cli/iiotools.py
index 15259d9..d381f8b 100644
--- a/bench/cli/iiotools.py
+++ b/bench/cli/iiotools.py
@@ -113,12 +113,21 @@ def capture_data(ctx, filename, device, channel, samples, props):
         for prop in props:
             if "=" not in prop:
                 raise ValueError(
-                    f"Invalid property: {prop}. Must be in the form key=value"
+                    f"Invalid property: {prop}. Must be in the form key=value\n",
+                    "or key=value,index for array properties",
                 )
             k, v = prop.split("=")
+            if "," in v:
+                v = v.split(",")
+                assert len(v) == 2, "Invalid array property"
+                v = v[0]
+                assert v.isdigit(), "Invalid array index property"
+                index = int(v[1])
+            else:
+                index = -1
             if v.isdigit():
                 v = int(v)
-            oprops[k] = v
+            oprops[k] = {"value": v, "index": index}
         props = oprops
 
     capture_iq_datafile(
diff --git a/bench/keysight/dwta/data_capture.py b/bench/keysight/dwta/data_capture.py
index 40e944f..a355bb6 100644
--- a/bench/keysight/dwta/data_capture.py
+++ b/bench/keysight/dwta/data_capture.py
@@ -32,11 +32,23 @@ def capture_iq_datafile(
     device = getattr(adi, device_name)(uri)
 
     # Assume kwargs are attributes to set
-    for key, value in kwargs.items():
-        print(f"Setting {key} to {value}")
+    for key, value_index in kwargs.items():
+        value = value_index["value"]
+        index = value_index["index"]
         if not hasattr(device, key):
             raise AttributeError(f"Device does not have attribute: {key}")
-        setattr(device, key, value)
+        if index == -1:
+            print(f"Setting {key} to {value}")
+            setattr(device, key, value)
+        else:
+            prop = getattr(device, key)
+            if prop is not list:
+                raise ValueError(f"Property {key} is not an array")
+            assert index < len(prop), f"Index out of range for {key}"
+            prop[index] = value
+            print(f"Setting {key}[{index}] to {value}")
+            setattr(device, key, prop)
+
 
     # Capture data
     device.rx_enabled_channels = [channel]