Merge branch 'main' into lecroy_socket_communication

scaaml/capture/scope/lecroy/lecroy_waveform.py

@@ -17,7 +17,7 @@
 import dataclasses
 import datetime
 import struct
-from typing import Union
+from typing import Tuple, Union
 
 import numpy as np
 
@@ -92,12 +92,18 @@ def get_double(self) -> float:
         return struct.unpack_from(f"{self._order:s}d", self.raw_data,
                                   self._ofs - 8)[0]
 
-    def get_unit(self):
-        """Parse a unit from self.raw_data.
+    def unpack_wave(self, amount: int) -> Tuple[int]:
+        """Parse the whole waveform from self.raw_data. Increments the ofset
+        `self._ofs` by the number of parsed bytes.
+
+        Args:
+          amount (int): Number of units (depending on self._unit).
+
+        Returns: A tuple of `amount` values.
         """
-        fmt = f"{self._order}{self._unit}"
+        fmt = f"{self._order}{amount}{self._unit}"
         size = struct.calcsize(fmt)
-        val = struct.unpack_from(fmt, self.raw_data, self._ofs)[0]
+        val = struct.unpack_from(fmt, self.raw_data, self._ofs)
         self._ofs += size
         return val
 
@@ -276,18 +282,14 @@ def parse(self):
         # Receive wave1
         block_len = self._wave_description.wave1_len
         if block_len > 0:
-            tmp = []
-            for _ in range(self._wave_description.wave_array_count):
-                tmp.append(self.get_unit())
-            self.d["wave1"] = np.array(tmp)
+            self.d["wave1"] = np.array(
+                self.unpack_wave(self._wave_description.wave_array_count))
 
         # Receive wave2
         block_len = self._wave_description.wave2_len
         if block_len > 0:
-            tmp = []
-            for _ in range(self._wave_description.wave_array_count):
-                tmp.append(self.get_unit())
-            self.d["wave2"] = np.array(tmp)
+            self.d["wave2"] = np.array(
+                self.unpack_wave(self._wave_description.wave_array_count))
         assert self._ofs == len(self.raw_data)
 
 

tests/capture/scope/lecroy/test_lecroy_waveform.py

@@ -0,0 +1,44 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Test LecroyWaveform."""
+
+import numpy as np
+
+from scaaml.capture.scope.lecroy.lecroy_waveform import LecroyWaveform
+
+
+def test_parse_waveform():
+    with open("tests/capture/scope/lecroy/raw_waveform.bin", "rb") as f:
+        raw_waveform = f.read()
+    lw = LecroyWaveform(raw_waveform)
+
+    assert lw._unit == "h"  # reading word values
+
+    expected = np.load("tests/capture/scope/lecroy/wave_1.npy",
+                       allow_pickle=False)
+
+    assert all(lw.get_wave1() == expected)
+
+
+def test_parse_waveform_byte():
+    with open("tests/capture/scope/lecroy/raw_waveform_byte.bin", "rb") as f:
+        raw_waveform = f.read()
+    lw = LecroyWaveform(raw_waveform)
+
+    assert lw._unit == "b"  # reading byte values
+
+    expected = np.load("tests/capture/scope/lecroy/wave_1_byte.npy",
+                       allow_pickle=False)
+
+    assert all(lw.get_wave1() == expected)