driver: SigrokDriver add blocking capture methods & misc improvements

This adds blocking capture methods where capture time or number of
samples are specified in advance.

Some of the internals are refactored to allow for code re-use
and API documentation is added to clarify the usage of the sigrok driver.

Signed-off-by: Felix Zwettler <[email protected]>

doc/configuration.rst

@@ -2710,6 +2710,9 @@ Arguments:
 The driver can be used in test cases by calling its ``capture()``, ``stop()``
 and ``analyze()`` methods.
 
+To capture for a certain predetermined amount of time or number of samples,
+methods ``capture_for_time()`` and ``capture_samples()`` can be used.
+
 SigrokPowerDriver
 ~~~~~~~~~~~~~~~~~
 The :any:`SigrokPowerDriver` uses a `SigrokUSBSerialDevice`_ resource to

labgrid/driver/sigrokdriver.py

@@ -92,16 +92,38 @@ def _get_sigrok_prefix(self):
         return self.sigrok.command_prefix + prefix
 
     @Driver.check_active
-    @step(title='call', args=['args'])
+    @step(title='call_with_driver', args=['args'])
     def _call_with_driver(self, *args):
         combined = self._get_sigrok_prefix() + list(args)
         self.logger.debug("Combined command: %s", " ".join(combined))
         self._process = subprocess.Popen(
             combined,
             stdout=subprocess.PIPE,
             stdin=subprocess.PIPE,
-            stderr=subprocess.PIPE
+            stderr=subprocess.PIPE,
+            text=True
+        )
+
+    @Driver.check_active
+    @step(title='call_with_driver_blocking', args=['args'])
+    def _call_with_driver_blocking(self, *args, log_output=False):
+        combined = self._get_sigrok_prefix() + list(args)
+        self.logger.debug("Combined command: %s", " ".join(combined))
+        process = subprocess.Popen(
+            combined,
+            stdout=subprocess.PIPE,
+            stdin=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True
         )
+        stdout, stderr = process.communicate()
+        if log_output:
+            self.logger.debug("stdout: %s", stdout)
+            self.logger.debug("stderr: %s", stderr)
+        if process.returncode != 0:
+            raise OSError
+        return stdout, stderr
+
 
     @Driver.check_active
     @step(title='call', args=['args'])
@@ -133,6 +155,19 @@ class SigrokDriver(SigrokCommon):
 
     @Driver.check_active
     def capture(self, filename, samplerate="200k"):
+        """
+        Starts to capture samples, in the background.
+
+        Args:
+            filename: the path to the file where the capture is saved.
+            samplerate: the sample-rate of the capture
+
+        Raises:
+            RuntimeError() if a capture is already running.
+        """
+        if self._running:
+            raise RuntimeError("capture is already running")
+
         self._filename = filename
         self._basename = os.path.basename(self._filename)
         self.logger.debug(
@@ -145,6 +180,7 @@ def capture(self, filename, samplerate="200k"):
         filename = os.path.join(self._tmpdir, self._basename)
         cmd.append(filename)
         self._call_with_driver(*cmd)
+
         args = self.sigrok.command_prefix + ['test', '-e', filename]
 
         while subprocess.call(args):
@@ -162,29 +198,115 @@ def capture(self, filename, samplerate="200k"):
 
         self._running = True
 
+    @Driver.check_active
+    def capture_for_time(self, filename, time_ms, samplerate="200k"):
+        """
+        Captures samples for a specified time (ms).
+
+        Blocks while capturing.
+
+        Args:
+            filename: the path to the file where the capture is saved.
+            time: time (in ms) for capture duration
+            samplerate: the sample-rate of the capture
+
+        Returns:
+            The capture as a list containing dict's with fields "Time"
+            and the channel names
+
+        Raises:
+            OSError() if the subprocess returned with non-zero return-code
+        """
+        return self._capture_blocking(filename, ["--time", str(time_ms)], samplerate)
+
+    @Driver.check_active
+    def capture_samples(self, filename, samples, samplerate="200k"):
+        """
+        Captures a specified number of samples.
+
+        Blocks while capturing.
+
+        Args:
+            filename: the path to the file where the capture is saved.
+            samples: number of samples to capture
+            samplerate: the sample-rate of the capture
+
+        Returns:
+            The capture as a list containing dict's with fields "Time"
+            and the channel names
+        """
+        return self._capture_blocking(filename, ["--samples", str(samples)], samplerate)
+
     @Driver.check_active
     def stop(self):
-        assert self._running
+        """
+        Stops the capture and returns recorded samples.
+
+        Note that this method might block for several seconds because it needs
+        to wait for output, parse the capture file and prepare the list
+        containing the samples.
+
+        Returns:
+            The capture as a list containing dict's with fields "Time"
+            and the channel names
+
+        Raises:
+            RuntimeError() if capture has not been started
+        """
+        if not self._running:
+            raise RuntimeError("no capture started yet")
         self._running = False
-        fnames = ['time']
-        fnames.extend(self.sigrok.channels.split(','))
+
         csv_filename = f'{os.path.splitext(self._basename)[0]}.csv'
 
         # sigrok-cli can be quit through any keypress
         stdout, stderr = self._process.communicate(input="q")
         self.logger.debug("stdout: %s", stdout)
         self.logger.debug("stderr: %s", stderr)
 
-        # Convert from .sr to .csv
+        self._convert_sr_csv(os.path.join(self._tmpdir, self._basename),
+                             os.path.join(self._tmpdir, csv_filename))
+        self._transfer_tmp_file(csv_filename)
+        return self._process_csv(csv_filename)
+
+    def _capture_blocking(self, filename, capture_args, samplerate):
+        self._filename = filename
+        self._basename = os.path.basename(self._filename)
+        csv_filename = f'{os.path.splitext(self._basename)[0]}.csv'
+        self.logger.debug(
+            "Saving to: %s with basename: %s", self._filename, self._basename
+        )
+        cmd = [
+            "-l", "4", "--config", f"samplerate={samplerate}",
+            *capture_args, "-o"
+        ]
+        filename = os.path.join(self._tmpdir, self._basename)
+        cmd.append(filename)
+        self._call_with_driver_blocking(*cmd, log_output=True)
+
+        args = self.sigrok.command_prefix + ['test', '-e', filename]
+
+        while subprocess.call(args):
+            sleep(0.1)
+
+        self._convert_sr_csv(os.path.join(self._tmpdir, self._basename),
+                             os.path.join(self._tmpdir, csv_filename))
+        self._transfer_tmp_file(csv_filename)
+        return self._process_csv(csv_filename)
+
+    def _convert_sr_csv(self, file_path_sr, file_path_scv):
         cmd = [
             '-i',
-            os.path.join(self._tmpdir, self._basename), '-O', 'csv:time=true', '-o',
-            os.path.join(self._tmpdir, csv_filename)
+            file_path_sr, '-O', 'csv:time=true', '-o',
+            file_path_scv
         ]
         self._call(*cmd)
         stdout, stderr = self._process.communicate()
         self.logger.debug("stdout: %s", stdout)
         self.logger.debug("stderr: %s", stderr)
+
+
+    def _transfer_tmp_file(self, csv_filename):
         if isinstance(self.sigrok, NetworkSigrokUSBDevice):
             subprocess.call([
                 'scp', f'{self.sigrok.host}:{os.path.join(self._tmpdir, self._basename)}',
@@ -201,16 +323,23 @@ def stop(self):
                             stdin=subprocess.DEVNULL,
                             stdout=subprocess.DEVNULL,
                             stderr=subprocess.DEVNULL)
+        else:
+            shutil.copyfile(
+                os.path.join(self._tmpdir, self._basename), self._filename
+            )
+
+    def _process_csv(self, csv_filename):
+        fnames = ['time']
+        fnames.extend(self.sigrok.channels.split(','))
+
+        if isinstance(self.sigrok, NetworkSigrokUSBDevice):
             with open(os.path.join(self._local_tmpdir,
                                    csv_filename)) as csv_file:
                 # skip first 5 lines of the csv output, contains metadata and fieldnames
                 for _ in range(0, 5):
                     next(csv_file)
                 return list(csv.DictReader(csv_file, fieldnames=fnames))
         else:
-            shutil.copyfile(
-                os.path.join(self._tmpdir, self._basename), self._filename
-            )
             with open(os.path.join(self._tmpdir, csv_filename)) as csv_file:
                 # skip first 5 lines of the csv output, contains metadata and fieldnames
                 for _ in range(0, 5):
@@ -219,6 +348,15 @@ def stop(self):
 
     @Driver.check_active
     def analyze(self, args, filename=None):
+        """
+        Analyzes captured data through `sigrok-cli`'s command line interface
+
+        Args:
+            args: args to `sigrok-cli`
+
+        Returns:
+            A dictionary containing the matched groups.
+        """
         annotation_regex = re.compile(r'(?P<startnum>\d+)-(?P<endnum>\d+) (?P<decoder>[\w\-]+): (?P<annotation>[\w\-]+): (?P<data>".*)')  # pylint: disable=line-too-long
         if not filename and self._filename:
             filename = os.path.join(self._tmpdir, self._basename)
@@ -408,12 +546,15 @@ def stop(self):
 
         Raises:
             RuntimeError() if capture has not been started
+            OSError() if the subprocess returned with non-zero return-code
         """
         if not self._running:
             raise RuntimeError("no capture started yet")
         while not self._timeout.expired:
             if self._process.poll() is not None:
                 # process has finished. no need to wait for the timeout
+                if self._process.returncode != 0:
+                    raise OSError
                 break
             time.sleep(0.1)
         else:

tests/test_sigrok.py

@@ -11,6 +11,8 @@
 pytestmark = pytest.mark.skipif(not which("sigrok-cli"),
                               reason="sigrok not available")
 
+VENDOR_ID = "0925"
+PRODUCT_ID = "3881"
 
 def test_sigrok_resource(target):
     r = SigrokUSBDevice(target, name=None, match={"sys_name": "1-12"}, driver='fx2lafw', channels="D0,D1")
@@ -23,19 +25,45 @@ def test_sigrok_driver(target):
 
 
 @pytest.mark.sigrokusb
-def test_sigrok_usb_driver(target, tmpdir):
-    r = SigrokUSBDevice(target, name=None, match={"ID_MODEL_ID": "3881", "ID_VENDOR_ID": "0925"}, driver='fx2lafw', channels="D0,D1")
+def test_sigrok_usb_driver_capture(target, tmpdir):
+    r = SigrokUSBDevice(target, name=None, match={"ID_MODEL_ID": PRODUCT_ID, "ID_VENDOR_ID": VENDOR_ID}, driver='fx2lafw', channels="D0,D1")
     d = SigrokDriver(target, name=None)
     target.activate(d)
     record = tmpdir.join("output.sr")
     d.capture(record)
     sleep(5)
     samples = d.stop()
     assert os.path.getsize(record) > 0
-    assert samples != None
+    assert samples is not None
     assert list(samples[0].keys()) == ['time', 'D0', 'D1']
     assert list(samples[-1].keys()) == ['time', 'D0', 'D1']
 
+@pytest.mark.sigrokusb
+def test_sigrok_usb_driver_blocking_samples(target, tmpdir):
+    r = SigrokUSBDevice(target, name=None, match={"ID_MODEL_ID": PRODUCT_ID, "ID_VENDOR_ID": VENDOR_ID}, driver='fx2lafw', channels="D0,D1")
+    d = SigrokDriver(target, name=None)
+    target.activate(d)
+    record = tmpdir.join("output.sr")
+    samples = d.capture_samples(record, 100)
+    assert os.path.getsize(record) > 0
+    assert samples is not None
+    assert len(samples) == 100
+
+
+@pytest.mark.sigrokusb
+def test_sigrok_usb_driver_blocking_time(target, tmpdir):
+    r = SigrokUSBDevice(target, name=None, match={"ID_MODEL_ID": PRODUCT_ID, "ID_VENDOR_ID": VENDOR_ID}, driver='fx2lafw', channels="D0,D1")
+    d = SigrokDriver(target, name=None)
+    target.activate(d)
+    record = tmpdir.join("output.sr")
+    samples = d.capture_for_time(record, 101) # sigrok-cli captures 5ms less than specified.
+    assert os.path.getsize(record) > 0
+    assert samples is not None
+    assert list(samples[0].keys()) == ['time', 'D0', 'D1']
+    assert list(samples[-1].keys()) == ['time', 'D0', 'D1']
+    time = float(samples[-1]['time']) - float(samples[0]['time'])
+    assert time >= 100_000
+
 def test_sigrok_power_driver(target):
     r = SigrokUSBSerialDevice(target, name=None, driver='manson-hcs-3xxx')
     r.avail = True