diff --git a/virl2_client/models/cl_pyats.py b/virl2_client/models/cl_pyats.py index 7886b04..d87cc0c 100644 --- a/virl2_client/models/cl_pyats.py +++ b/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: diff --git a/virl2_client/models/node.py b/virl2_client/models/node.py index c779344..0964ff7 100644 --- a/virl2_client/models/node.py +++ b/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