diff --git a/scaaml/capture/scope/lecroy/lecroy_waveform.py b/scaaml/capture/scope/lecroy/lecroy_waveform.py index b8fcc55d..70aa5653 100644 --- a/scaaml/capture/scope/lecroy/lecroy_waveform.py +++ b/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) diff --git a/tests/capture/scope/lecroy/raw_waveform.bin b/tests/capture/scope/lecroy/raw_waveform.bin new file mode 100644 index 00000000..e151325e Binary files /dev/null and b/tests/capture/scope/lecroy/raw_waveform.bin differ diff --git a/tests/capture/scope/lecroy/raw_waveform_byte.bin b/tests/capture/scope/lecroy/raw_waveform_byte.bin new file mode 100644 index 00000000..a4c0ba8f Binary files /dev/null and b/tests/capture/scope/lecroy/raw_waveform_byte.bin differ diff --git a/tests/capture/scope/lecroy/test_lecroy_waveform.py b/tests/capture/scope/lecroy/test_lecroy_waveform.py new file mode 100644 index 00000000..63a1b79f --- /dev/null +++ b/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) diff --git a/tests/capture/scope/lecroy/wave_1.npy b/tests/capture/scope/lecroy/wave_1.npy new file mode 100644 index 00000000..50135c19 Binary files /dev/null and b/tests/capture/scope/lecroy/wave_1.npy differ diff --git a/tests/capture/scope/lecroy/wave_1_byte.npy b/tests/capture/scope/lecroy/wave_1_byte.npy new file mode 100644 index 00000000..466e1604 Binary files /dev/null and b/tests/capture/scope/lecroy/wave_1_byte.npy differ