fix: validate waveform tables before device writes

pslab/instrument/waveform_generator.py

@@ -253,6 +253,11 @@ def load_function(
             The minimum and maximum x values between which to evaluate
             'function'. Should typically correspond to one period. Omit if
             'function' is 'sine' or 'tria'.
+
+        Raises
+        ------
+        ValueError
+            If the function does not return 512 finite real voltage values.
         """
         if function == "sine":
 
@@ -261,23 +266,21 @@ def sine(x):
 
             function = sine
             span = [0, 2 * np.pi]
-            self._channels[channel].wavetype = "sine"
+            mode = "sine"
         elif function == "tria":
 
             def tria(x):
                 return AnalogOutput.RANGE[1] * (abs(x % 4 - 2) - 1)
 
             function = tria
             span = [0, 4]
-            self._channels[channel].wavetype = "tria"
+            mode = "tria"
         else:
-            self._channels[channel].wavetype = "custom"
+            mode = "custom"
 
         x = np.arange(span[0], span[1], (span[1] - span[0]) / 512)
         y = function(x)
-        self._load_table(
-            channel=channel, points=y, mode=self._channels[channel].wavetype
-        )
+        self._load_table(channel=channel, points=y, mode=mode)
 
     def load_table(self, channel: str, points: np.ndarray):
         """Load a custom waveform as a table.
@@ -289,13 +292,38 @@ def load_table(self, channel: str, points: np.ndarray):
         points : np.ndarray
             Array of voltage values which make up the waveform. Array length
             must be 512. Values outside the range -3.3 V to 3.3 V will be
-            clipped.
+            clipped. All values must be finite real numbers. Complex values and
+            booleans are rejected.
+
+        Raises
+        ------
+        ValueError
+            If the table is not a one-dimensional array of 512 finite real
+            voltage values. Invalid tables are rejected before device writes.
         """
         self._load_table(channel, points, "custom")
 
     def _load_table(self, channel, points, mode="custom"):
-        self._channels[channel].wavetype = mode
+        points = np.asarray(points)
+        if points.ndim != 1 or points.size != self._HIGHRES_TABLE_SIZE:
+            raise ValueError("Waveform table must contain exactly 512 values.")
+        if points.dtype.kind not in "iufOUS" or (
+            points.dtype.kind == "O"
+            and any(isinstance(value, (bool, np.bool_)) for value in points)
+        ):
+            raise ValueError("Waveform voltages must be finite real values.")
+        if points.dtype.kind in "OUS":
+            try:
+                points = points.astype(float)
+            except (TypeError, ValueError, OverflowError) as error:
+                raise ValueError(
+                    "Waveform voltages must be finite real values."
+                ) from error
+        if not np.isfinite(points).all():
+            raise ValueError("Waveform voltages must be finite real values.")
+
         self._channels[channel].waveform_table = points
+        self._channels[channel].wavetype = mode
         logger.info(f"Reloaded waveform table for {channel}: {mode}.")
         self._device.send_byte(CP.WAVEGEN)
 

tests/test_waveform_table_validation.py

@@ -0,0 +1,82 @@
+"""Hardware-independent validation of waveform uploads and command framing."""
+
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+
+import pslab.protocol as CP
+from pslab.connection import ConnectionHandler
+from pslab.instrument.waveform_generator import WaveformGenerator
+
+
+@pytest.mark.parametrize("loader", ["table", "function"])
+@pytest.mark.parametrize(
+    "points",
+    [
+        np.zeros(0),
+        np.zeros(1),
+        np.zeros(511),
+        np.zeros(513),
+        np.zeros((512, 1)),
+        np.zeros((2, 256)),
+        np.full(512, np.nan),
+        np.full(512, np.inf),
+        np.full(512, -np.inf),
+        np.full(512, 1j),
+        np.full(512, "not-a-voltage"),
+        np.full(512, None, dtype=object),
+        np.full(512, {}, dtype=object),
+        np.resize(np.array([False, True]), 512),
+        np.full(512, True, dtype=object),
+        np.array([np.bool_(False)] + [0.5] * 511, dtype=object),
+    ],
+)
+def test_invalid_table_does_not_write_or_change_channel(points, loader):
+    device = Mock(spec=ConnectionHandler)
+    generator = WaveformGenerator(device)
+    channel = generator._channels["SI1"]
+    previous_table = channel.waveform_table.copy()
+    with pytest.raises(ValueError):
+        if loader == "table":
+            generator.load_table("SI1", points)
+        else:
+            generator.load_function("SI1", lambda x: points, [0, 1])
+    assert device.mock_calls == []
+    assert channel.wavetype == "sine"
+    assert channel.waveform_table == previous_table
+
+
+@pytest.mark.parametrize("channel", ["SI1", "SI2"])
+@pytest.mark.parametrize("dtype", [np.float64, np.int16, object])
+def test_valid_table_preserves_upload_framing(channel, dtype):
+    device = Mock(spec=ConnectionHandler)
+    generator = WaveformGenerator(device)
+    points = np.linspace(-5, 5, 512).astype(dtype)
+    original = points.copy()
+    generator.load_table(channel, points)
+    clipped = np.clip(np.asarray(points, dtype=float), -3.3, 3.3)
+    expected = np.round((clipped + 3.3) / 6.6 * 511).astype(np.int16).tolist()
+    assert [call.args[0] for call in device.send_int.call_args_list] == expected
+    commands = [call.args[0] for call in device.send_byte.call_args_list]
+    assert commands[:2] == [
+        CP.WAVEGEN,
+        CP.LOAD_WAVEFORM1 if channel == "SI1" else CP.LOAD_WAVEFORM2,
+    ]
+    assert (
+        commands[2:]
+        == np.round((clipped[::16] + 3.3) / 6.6 * 63).astype(np.int16).tolist()
+    )
+    device.get_ack.assert_called_once_with()
+    assert generator._channels[channel].wavetype == "custom"
+    np.testing.assert_array_equal(points, original)
+
+
+@pytest.mark.parametrize("function", ["sine", "tria"])
+def test_builtin_functions_keep_fixed_upload_sizes(function):
+    device = Mock(spec=ConnectionHandler)
+    generator = WaveformGenerator(device)
+    generator.load_function("SI1", function)
+    assert device.send_int.call_count == 512
+    assert device.send_byte.call_count == 34
+    assert generator._channels["SI1"].wavetype == function