Allow passing arbitrary PyATS parameters in run command

virl2_client/models/cl_pyats.py

@@ -21,7 +21,7 @@
 from __future__ import annotations
 
 import io
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 try:
     from pyats.topology.loader.base import TestbedFileLoader as _PyatsTFLoader
@@ -125,6 +125,7 @@ def _prepare_params(
         self,
         init_exec_commands: list[str] | None = None,
         init_config_commands: list[str] | None = None,
+        **params: str,
     ) -> dict:
         """
         Prepare a dictionary of optional parameters to be executed before a command.
@@ -135,10 +136,9 @@ def _prepare_params(
         :param init_config_commands: A list of config commands to be executed.
         :returns: A dictionary of optional parameters to be executed with a command.
         """
-        params = {}
-        if init_exec_commands:
+        if init_exec_commands is not None:
             params["init_exec_commands"] = init_exec_commands
-        if init_config_commands:
+        if init_config_commands is not None:
             params["init_config_commands"] = init_config_commands
         return params
 
@@ -149,6 +149,7 @@ def _execute_command(
         configure_mode: bool = False,
         init_exec_commands: list[str] | None = None,
         init_config_commands: list[str] | None = None,
+        **pyats_params: Any,
     ) -> str:
         """
         Execute a command on the device.
@@ -163,6 +164,7 @@ def _execute_command(
         :param init_config_commands: A list of config commands to be executed
             before the command. Default commands will be run if omitted.
             Pass an empty list to run no commands.
+        :param pyats_params: Additional PyATS call parameters
         :returns: The output from the device.
         :raises PyatsDeviceNotFound: If the device cannot be found.
         """
@@ -179,7 +181,9 @@ def _execute_command(
 
             pyats_device.connect(log_stdout=False, learn_hostname=True)
             self._connections.add(pyats_device)
-        params = self._prepare_params(init_exec_commands, init_config_commands)
+        params = self._prepare_params(
+            init_exec_commands, init_config_commands, **pyats_params
+        )
         if configure_mode:
             return pyats_device.configure(command, log_stdout=False, **params)
         else:
@@ -191,6 +195,7 @@ def run_command(
         command: str,
         init_exec_commands: list[str] | None = None,
         init_config_commands: list[str] | None = None,
+        **pyats_params: Any,
     ) -> str:
         """
         Run a command on the device in exec mode.
@@ -203,6 +208,7 @@ def run_command(
         :param init_config_commands: A list of config commands to be executed
             before the command. Default commands will be run if omitted.
             Pass an empty list to run no commands.
+        :param pyats_params: Additional PyATS call parameters
         :returns: The output from the device.
         """
         return self._execute_command(
@@ -211,6 +217,7 @@ def run_command(
             configure_mode=False,
             init_exec_commands=init_exec_commands,
             init_config_commands=init_config_commands,
+            **pyats_params,
         )
 
     def run_config_command(
@@ -219,6 +226,7 @@ def run_config_command(
         command: str,
         init_exec_commands: list[str] | None = None,
         init_config_commands: list[str] | None = None,
+        **pyats_params: Any,
     ) -> str:
         """
         Run a command on the device in configure mode. pyATS automatically handles the
@@ -232,6 +240,7 @@ def run_config_command(
         :param init_config_commands: A list of config commands to be executed
             before the command. Default commands will be run if omitted.
             Pass an empty list to run no commands.
+        :param pyats_params: Additional PyATS call parameters
         :returns: The output from the device.
         """
         return self._execute_command(
@@ -240,6 +249,7 @@ def run_config_command(
             configure_mode=True,
             init_exec_commands=init_exec_commands,
             init_config_commands=init_config_commands,
+            **pyats_params,
         )
 
     def cleanup(self) -> None:

virl2_client/models/node.py

@@ -860,25 +860,27 @@ def _remove_tag_on_server(self, tag) -> None:
         current.remove(tag)
         self._set_node_property("tags", current)
 
-    def run_pyats_command(self, command: str) -> str:
+    def run_pyats_command(self, command: str, **pyats_params: Any) -> str:
         """
         Run a pyATS command in exec mode on the node.
 
         :param command: The command to run (e.g. "show version").
+        :param pyats_params: Custom command dialog parameters for PyATS
         :returns: The output from the device.
         """
         label = self.label
-        return self._lab.pyats.run_command(label, command)
+        return self._lab.pyats.run_command(label, command, **pyats_params)
 
-    def run_pyats_config_command(self, command: str) -> str:
+    def run_pyats_config_command(self, command: str, **pyats_params: Any) -> str:
         """
         Run a pyATS command in config mode on the node.
 
         :param command: The command to run (e.g. "interface gi0").
+        :param pyats_params: Custom command dialog parameters for PyATS
         :returns: The output from the device.
         """
         label = self.label
-        return self._lab.pyats.run_config_command(label, command)
+        return self._lab.pyats.run_config_command(label, command, **pyats_params)
 
     @check_stale
     @locked