gists/python/state-cli.md

Implementation of a command line interface that provides different commands based on the state of a state machine.

import prompt_toolkit
from prompt_toolkit.styles import Style
from prompt_toolkit.patch_stdout import StdoutProxy, patch_stdout
import time
import threading
import logging


#
# StateCli
#


class StateCommandBase:
    """Base class for state-specific commands."""

    def __init__(self, cli, state_machine):
        self.cli = cli
        self.state_machine = state_machine


class StateCli:
    """
    Implementation of a CLI working in conjuction with a state machine.
    Depending on the state, the CLI provides different commands.
    """

    def __init__(self, state_machine, command_map):
        self.state_machine = state_machine

        self.state_commands = {
            key: value(self, state_machine) for key, value in command_map.items()
        }

        self.update_prompt()

        self._prompt_style = Style.from_dict(
            {
                "prompt": "ansiblue bold",
                "input": "",
            }
        )

        self.session = prompt_toolkit.PromptSession()

    def update_prompt(self):
        self.prompt = f"{self.state_machine.current_state}> "

    def _process_command(self, command, args):
        current_state = self.state_machine.current_state
        if current_state not in self.state_commands:
            print(f"No commands defined for state {current_state}")
            return

        state_cmd_obj = self.state_commands[current_state]
        if not hasattr(state_cmd_obj, command):
            print(f"Unknown command: '{command}'")
            self._do_help()
            return

        getattr(state_cmd_obj, command)(" ".join(args))

    def cmdloop(self):
        with patch_stdout():
            while True:
                try:
                    user_input = self.session.prompt(
                        self.prompt, style=self._prompt_style
                    )

                    parts = user_input.split()
                    if len(parts) == 0:
                        continue

                    command, args = parts[0], parts[1:]

                    if command == "help":
                        self._do_help()
                        continue
                    elif command == "exit":
                        break

                    self._process_command(command, args)
                except KeyboardInterrupt:  # ctrl+c
                    break
                except EOFError:  # ctrl+d
                    break

    def _do_help(self):
        current_state = self.state_machine.current_state
        print("=" * 50)
        print("Available commands in any state")
        print(f"    {'help':15} - Show this help")
        print(f"    {'exit':15} - Exit the application")

        print("=" * 50)
        print(f"Available commands in state {current_state}")

        if current_state not in self.state_commands:
            print(f"   No commands defined for state {current_state}")
            print("=" * 50)
            return

        state_cmd_obj = self.state_commands[current_state]
        commands = [
            name
            for name in dir(state_cmd_obj)
            if not name.startswith("_") and callable(getattr(state_cmd_obj, name))
        ]

        if not commands:
            print("    No commands available")
            print("=" * 50)
            return

        for cmd in sorted(commands):
            method = getattr(state_cmd_obj, cmd)
            if method.__doc__:
                doc = method.__doc__.strip()
                first_line = doc.split("\n")[0]
                print(f"    {cmd:15} - {first_line}")
            else:
                print(f"    {cmd:15} - No documentation available")

        print("=" * 50)


#
# StateMachine
#


class StateMachine:
    def __init__(self):
        self.current_state = "idle"

    def start(self):
        self.current_state = "running"

    def stop(self):
        self.current_state = "idle"


#
# Usage
#
 

class IdleCommands(StateCommandBase):
    def start(self, args):
        """Start the machine"""
        print("Starting the machine")
        self.state_machine.start()
        self.cli.update_prompt()


class RunningCommands(StateCommandBase):
    def _stop(self, args):
        """Stop the machine"""
        print("Stopping the the machine")
        self.state_machine.stop()
        self.cli.update_prompt()


def background_task():
    while True:
        time.sleep(2)
        logging.info("This is a log message from some background task")


def main():
    sm = StateMachine()
    cli = StateCli(sm, {"idle": IdleCommands, "running": RunningCommands})

    logging.basicConfig(
        format="%(asctime)s %(levelname)s %(message)s",
        datefmt="%H:%M:%S",
        level=logging.INFO,
        stream=StdoutProxy(),
    )

    thread = threading.Thread(target=background_task, daemon=True)
    thread.start()

    cli.cmdloop()


if __name__ == "__main__":
    main()