shell_utils

drunc.utils.shell_utils

Shell utilities for DRUNC.

Classes

CommandLike

Bases: Protocol

Protocol for command-like objects.

ControllerDriverProtocol

Bases: Protocol

Protocol for controller driver objects.

Methods:

describe_fsm()

Describe the FSM.

Returns:

Name	Type	Description
DescribeFSMReplyLike	DescribeFSMReplyLike	The FSM description.

Source code in drunc/utils/shell_utils.py

def describe_fsm(self) -> DescribeFSMReplyLike:
    """Describe the FSM.

    Returns:
        DescribeFSMReplyLike: The FSM description.
    """
    ...

status()

Get the current status.

Returns:

Name	Type	Description
StatusReplyLike	StatusReplyLike	The current status.

Source code in drunc/utils/shell_utils.py

def status(self) -> StatusReplyLike:
    """Get the current status.

    Returns:
        StatusReplyLike: The current status.
    """
    ...

DecodedResponse(name, token, flag, data=None, children=None)

Decoded response object.

Warning: This should be kept in sync with druncschema/request_response.proto/Response class

Initialize a DecodedResponse.

Parameters:

Name	Type	Description	Default
name	str	The name of the response.	required
token	Token	The token associated with the response.	required
flag	object	The response flag.	required
data	object | None	The response data. Defaults to None.	None
children	list[DecodedResponse] | None	Child responses. Defaults to None.	None

Source code in drunc/utils/shell_utils.py

def __init__(
    self,
    name: str,
    token: Token,
    flag: object,
    data: object | None = None,
    children: list["DecodedResponse"] | None = None,
) -> None:
    """Initialize a DecodedResponse.

    Args:
        name: The name of the response.
        token: The token associated with the response.
        flag: The response flag.
        data: The response data. Defaults to None.
        children: Child responses. Defaults to None.
    """
    self.name = name
    self.token = token
    self.flag = flag
    self.data = data
    if children is None:
        self.children = []
    else:
        self.children = children

Methods:

__str__()

Return string representation of the DecodedResponse.

Returns:

Name	Type	Description
str	str	The string representation.

Source code in drunc/utils/shell_utils.py

def __str__(self) -> str:
    """Return string representation of the DecodedResponse.

    Returns:
        str: The string representation.
    """
    return DecodedResponse.to_string(self)

to_string(obj, prefix='') staticmethod

Convert a DecodedResponse to a string representation.

Parameters:

Name	Type	Description	Default
obj	DecodedResponse	The DecodedResponse to convert.	required
prefix	str	A prefix to add to the string. Defaults to empty string.	''

Returns:

Name	Type	Description
str	str	The string representation of the response.

Source code in drunc/utils/shell_utils.py

@staticmethod
def to_string(obj: "DecodedResponse", prefix: str = "") -> str:
    """Convert a DecodedResponse to a string representation.

    Args:
        obj: The DecodedResponse to convert.
        prefix: A prefix to add to the string. Defaults to empty string.

    Returns:
        str: The string representation of the response.
    """
    text = (
        f"{prefix} {obj.name} -> response flag={obj.flag} type={type(obj.data)}\n"
    )
    for v in obj.children:
        if v is None:
            continue
        text += DecodedResponse.to_string(v, prefix + "  ")
    return text

DescribeFSMReplyLike

Bases: Protocol

Protocol for describe FSM reply-like objects.

FSMDescriptionLike

Bases: Protocol

Protocol for FSM description-like objects.

InterruptedCommand(message='An error occurred in Drunc.', grpc_error_code=None, details=None, **detail_kwargs)

Bases: DruncShellException

Exception thrown to interrupt a shell command without a full stack trace.

Source code in drunc/exceptions.py

def __init__(
    self,
    message: str = "An error occurred in Drunc.",
    grpc_error_code=None,
    details=None,  # optional rich error detail
    **detail_kwargs,
):
    super().__init__(message)

    if message is not None:
        self.message = message

    if grpc_error_code is None:
        grpc_error_code = getattr(self, "grpc_error_code", code_pb2.INTERNAL)

    if details is not None:
        self.details = details

    self.detail_kwargs = detail_kwargs

SequenceLike

Bases: Protocol

Protocol for sequence-like objects.

ShellContext(*args, **kwargs)

Base class for shell contexts.

Initialize the shell context.

Parameters:

Name	Type	Description	Default
*args	object	Additional positional arguments.	()
**kwargs	object	Additional keyword arguments.	{}

Source code in drunc/utils/shell_utils.py

def __init__(self, *args: object, **kwargs: object) -> None:
    """Initialize the shell context.

    Args:
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.
    """
    log = get_logger("utils.ShellContext")
    self.dynamic_commands: set[str] = set()
    try:
        self.reset(*args, **kwargs)
    except Exception as e:
        log.exception(e)
        exit(1)

Methods:

create_drivers(**kwargs) abstractmethod

Create drivers for the context.

Parameters:

Name	Type	Description	Default
**kwargs	object	Additional keyword arguments.	{}

Returns:

Type	Description
MutableMapping[str, object]	MutableMapping[str, object]: A mapping of driver names to driver objects.

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def create_drivers(self, **kwargs: object) -> MutableMapping[str, object]:
    """Create drivers for the context.

    Args:
        **kwargs: Additional keyword arguments.

    Returns:
        MutableMapping[str, object]: A mapping of driver names to driver objects.
    """
    pass

create_token(**kwargs) abstractmethod

Create a token for the context.

Parameters:

Name	Type	Description	Default
**kwargs	object	Additional keyword arguments.	{}

Returns:

Name	Type	Description
Token	Token	A token object.

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def create_token(self, **kwargs: object) -> Token:
    """Create a token for the context.

    Args:
        **kwargs: Additional keyword arguments.

    Returns:
        Token: A token object.
    """
    pass

delete_driver(name)

Delete a driver from the context.

Parameters:

Name	Type	Description	Default
name	str	The name of the driver to delete.	required

Source code in drunc/utils/shell_utils.py

def delete_driver(self, name: str) -> None:
    """Delete a driver from the context.

    Args:
        name: The name of the driver to delete.
    """
    log = get_logger("utils.ShellContext")
    if name in self._drivers:
        log.info(f"You will not be able to issue commands to the {name} anymore.")
        del self._drivers[name]
        log.info(f"{name.capitalize()} driver has been deleted.")

get_driver(name=None, quiet_fail=False)

Get a driver from the context.

Parameters:

Name	Type	Description	Default
name	str | None	The name of the driver. If None, returns the only driver if there is exactly one.	None
quiet_fail	bool	If True, return None on failure instead of raising an exception.	False

Returns:

Name	Type	Description
object	object	The driver object, or None if quiet_fail is True and the driver is not found.

Raises:

Type	Description
DruncShellException	If there are multiple drivers and no name is specified.
SystemExit	If the driver is not found and quiet_fail is False.

Source code in drunc/utils/shell_utils.py

def get_driver(self, name: str | None = None, quiet_fail: bool = False) -> object:
    """Get a driver from the context.

    Args:
        name: The name of the driver. If None, returns the only driver if there is exactly one.
        quiet_fail: If True, return None on failure instead of raising an exception.

    Returns:
        object: The driver object, or None if quiet_fail is True and the driver is not found.

    Raises:
        DruncShellException: If there are multiple drivers and no name is specified.
        SystemExit: If the driver is not found and quiet_fail is False.
    """
    try:
        if name:
            return self._drivers[name]
        elif len(self._drivers) > 1:
            raise DruncShellException("More than one driver in this context")
        return list(self._drivers.values())[0]
    except KeyError:
        if quiet_fail:
            return None
        log = get_logger("utils.ShellContext")
        log.exception(
            "Controller-specific commands cannot be sent until the session is booted"
        )
        log.debug(f"Drivers available are {self._drivers.keys()}")
        raise SystemExit(
            1
        )  # used to avoid having to catch multiple Attribute errors when this function gets called

get_token()

Get the token from the context.

Returns:

Name	Type	Description
Token	Token	The token object.

Source code in drunc/utils/shell_utils.py

def get_token(self) -> Token:
    """Get the token from the context.

    Returns:
        Token: The token object.
    """
    return self._token

has_driver(name)

Check if a driver exists in the context.

Parameters:

Name	Type	Description	Default
name	str	The name of the driver.	required

Returns:

Name	Type	Description
bool	bool	True if the driver exists, False otherwise.

Source code in drunc/utils/shell_utils.py

def has_driver(self, name: str) -> bool:
    """Check if a driver exists in the context.

    Args:
        name: The name of the driver.

    Returns:
        bool: True if the driver exists, False otherwise.
    """
    return name in self._drivers

print(*args, **kwargs)

Print to the console.

Parameters:

Name	Type	Description	Default
*args	object	Positional arguments to pass to the console.	()
**kwargs	object	Keyword arguments to pass to the console.	{}

Source code in drunc/utils/shell_utils.py

def print(self, *args: object, **kwargs: object) -> None:
    """Print to the console.

    Args:
        *args: Positional arguments to pass to the console.
        **kwargs: Keyword arguments to pass to the console.
    """
    self._console.print(*args, **kwargs)  # type: ignore[arg-type]

print_status_summary()

Print a summary of the FSM status and available transitions.

Source code in drunc/utils/shell_utils.py

def print_status_summary(self) -> None:
    """Print a summary of the FSM status and available transitions."""
    log = get_logger("utils.ShellContext")
    controller = cast(ControllerDriverProtocol, self.get_driver("controller"))
    status = controller.status().status
    describe_fsm = controller.describe_fsm().description
    current_state = status.state
    if status.in_error:
        log.error(
            f"[red] FSM is in error ({status})[/red], not currently accepting new commands."
        )
    else:
        available_actions = [
            command.name.replace("_", "-") for command in describe_fsm.commands
        ]
        available_sequences = [
            seq.id.replace("_", "-") for seq in describe_fsm.sequences
        ]

        log.info(
            f"Current FSM status is [green]{current_state}[/green]. Available transitions are [green]{'[/green], [green]'.join(available_actions)}[/green]. Available sequence commands are [green]{'[/green], [green]'.join(available_sequences)}[/green]."
        )

reset(**kwargs) abstractmethod

Reset the shell context.

Parameters:

Name	Type	Description	Default
**kwargs	object	Additional keyword arguments.	{}

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def reset(self, **kwargs: object) -> None:
    """Reset the shell context.

    Args:
        **kwargs: Additional keyword arguments.
    """
    pass

rule(*args, **kwargs)

Print a rule to the console.

Parameters:

Name	Type	Description	Default
*args	object	Positional arguments to pass to the console.	()
**kwargs	object	Keyword arguments to pass to the console.	{}

Source code in drunc/utils/shell_utils.py

def rule(self, *args: object, **kwargs: object) -> None:
    """Print a rule to the console.

    Args:
        *args: Positional arguments to pass to the console.
        **kwargs: Keyword arguments to pass to the console.
    """
    self._console.rule(*args, **kwargs)  # type: ignore[arg-type]

set_driver(name, driver)

Set a driver in the context.

Parameters:

Name	Type	Description	Default
name	str	The name of the driver.	required
driver	object	The driver object.	required

Raises:

Type	Description
DruncShellException	If a driver with the same name already exists.

Source code in drunc/utils/shell_utils.py

def set_driver(self, name: str, driver: object) -> None:
    """Set a driver in the context.

    Args:
        name: The name of the driver.
        driver: The driver object.

    Raises:
        DruncShellException: If a driver with the same name already exists.
    """
    if name in self._drivers:
        raise DruncShellException(f"Driver {name} already present in this context")
    self._drivers[name] = driver

terminate() abstractmethod

Terminate the shell context.

Source code in drunc/utils/shell_utils.py

@abc.abstractmethod
def terminate(self) -> None:
    """Terminate the shell context."""
    pass

StatusLike

Bases: Protocol

Protocol for status-like objects.

StatusReplyLike

Bases: Protocol

Protocol for status reply-like objects.

Functions:

add_traceback_flag()

Add a traceback flag to a command.

Returns:

Name	Type	Description
Callable	Callable[[Callable[P, R]], Callable[P, R]]	A decorator that adds the traceback flag.

Source code in drunc/utils/shell_utils.py

def add_traceback_flag() -> Callable[[Callable[P, R]], Callable[P, R]]:
    """Add a traceback flag to a command.

    Returns:
        Callable: A decorator that adds the traceback flag.
    """

    def wrapper(f0: Callable[P, R]) -> Callable[P, R]:
        f1 = click.option(
            "-t/-nt",
            "--traceback/--no-traceback",
            default=None,
            help="Print full exception traceback",
        )(f0)
        return f1

    return wrapper

create_dummy_token_from_uname()

Create a dummy token from the current username.

Returns:

Name	Type	Description
Token	Token	A dummy token with the current username.

Source code in drunc/utils/shell_utils.py

def create_dummy_token_from_uname() -> Token:
    """Create a dummy token from the current username.

    Returns:
        Token: A dummy token with the current username.
    """
    user = getpass.getuser()
    return (
        Token(  # fake token, but should be figured out from the environment/authoriser
            token=f"{user}-token", user_name=user
        )
    )