Merge pull request #1345 from flxzt/sigrok-driver-improv

driver/sigrok: add blocking capture methods, misc improvements to sigrok driver

Author: Emantor

doc/configuration.rst

@@ -3127,6 +3127,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

@@ -94,17 +94,39 @@ 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'])
     def _call(self, *args):
@@ -137,8 +159,22 @@ 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._validate_save_dir_exists()
         self.logger.debug(
             "Saving to: %s with basename: %s", self._filename, self._basename
         )
@@ -149,6 +185,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):
@@ -161,60 +198,207 @@ def capture(self, filename, samplerate="200k"):
                     self.logger.debug("sigrok-cli call terminated prematurely with non-zero return-code")
                     self.logger.debug("stdout: %s", stdout)
                     self.logger.debug("stderr: %s", stderr)
-                    raise ExecutionError(f"sigrok-cli call terminated prematurely with return-code '{ret}'.")
+                    raise ExecutionError(f"sigrok-cli call terminated prematurely, return-code '{ret}'.")
             sleep(0.1)
 
         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=b"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)
+        self._validate_save_dir_exists()
+        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 _validate_save_dir_exists(self):
+        if not os.path.exists(os.path.dirname(self._filename)):
+            raise ExecutionError(
+                f"Won't be able to save capture file to target directory '{os.path.dirname(self._filename)}', does not exist"
+            )
+
+    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)}',
-                os.path.abspath(self._filename)
-            ],
-                            stdin=subprocess.DEVNULL,
-                            stdout=subprocess.DEVNULL,
-                            stderr=subprocess.DEVNULL)
+            sr_remote_path = os.path.join(self._tmpdir, self._basename)
+            sr_local_path = os.path.abspath(self._filename)
+            csv_remote_path = os.path.join(self._tmpdir, csv_filename)
+            csv_local_path = os.path.join(self._local_tmpdir, csv_filename)
+            self.logger.debug(
+                "Transferring sigrok capture file from remote path '%s' to local path '%s'",
+                sr_remote_path,
+                sr_local_path
+            )
+            # TODO: use NetworkResource ssh base options here too
+            ret = subprocess.run([
+                    'scp',
+                    "-o",
+                    "UserKnownHostsFile=/dev/null",
+                    "-o",
+                    "StrictHostKeyChecking=no",
+                    f'{self.sigrok.host}:{sr_remote_path}',
+                    sr_local_path,
+                ],
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT
+            )
+            if ret.returncode != 0:
+                self.logger.error(
+                    "Transferring Sigrok capture file failed, return-code '%i : %s'",
+                    ret.returncode,
+                    ret.stdout.decode()
+                )
+                raise ExecutionError(
+                    f"Transferring Sigrok Capture file failed, return-code '{ret.returncode}'",
+                    stdout=ret.stdout.decode().split("\n")
+                )
+
             # get csv from remote host
-            subprocess.call([
-                'scp', f'{self.sigrok.host}:{os.path.join(self._tmpdir, csv_filename)}',
-                os.path.join(self._local_tmpdir, csv_filename)
-            ],
-                            stdin=subprocess.DEVNULL,
-                            stdout=subprocess.DEVNULL,
-                            stderr=subprocess.DEVNULL)
+            self.logger.debug(
+                "Transferring CSV file from remote path '%s' to local path '%s'",
+                csv_remote_path,
+                csv_local_path
+            )
+            # TODO: use NetworkResource ssh base options here too
+            ret = subprocess.run([
+                    'scp',
+                    "-o",
+                    "UserKnownHostsFile=/dev/null",
+                    "-o",
+                    "StrictHostKeyChecking=no",
+                    f'{self.sigrok.host}:{csv_remote_path}',
+                    csv_local_path
+                ],
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT)
+            if ret.returncode != 0:
+                self.logger.error(
+                    "Transferring Sigrok CSV file failed, return-code '%i : %s'",
+                    ret.returncode,
+                    ret.stdout.decode()
+                )
+                raise ExecutionError(
+                    f"Transferring Sigrok CSV file failed, return-code '{ret.returncode}'",
+                    stdout=ret.stdout.decode().split("\n")
+                )
+        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):
@@ -223,6 +407,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)
@@ -412,12 +605,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: