From 23ef39a2d99875c2adaabf562f96bbab5a733464 Mon Sep 17 00:00:00 2001 From: a-kenji Date: Thu, 30 May 2024 19:51:53 +0200 Subject: [PATCH] clan: add dynamic completions Add dynamic completion scaffolding to the clan `cli`. Also add a dynamic completion mechanism for machines for commands that have machines as their sole argument. More intricate dynamic completions will be implemented in follow up PR's. --- pkgs/clan-cli/clan_cli/__init__.py | 1 + pkgs/clan-cli/clan_cli/backups/create.py | 5 +- pkgs/clan-cli/clan_cli/backups/list.py | 4 +- pkgs/clan-cli/clan_cli/completions.py | 81 ++++++++++++++++++++++ pkgs/clan-cli/clan_cli/facts/check.py | 5 +- pkgs/clan-cli/clan_cli/facts/generate.py | 5 +- pkgs/clan-cli/clan_cli/facts/list.py | 5 +- pkgs/clan-cli/clan_cli/facts/upload.py | 5 +- pkgs/clan-cli/clan_cli/flash.py | 5 +- pkgs/clan-cli/clan_cli/machines/delete.py | 5 +- pkgs/clan-cli/clan_cli/machines/install.py | 6 +- pkgs/clan-cli/clan_cli/machines/update.py | 6 +- 12 files changed, 123 insertions(+), 10 deletions(-) create mode 100644 pkgs/clan-cli/clan_cli/completions.py diff --git a/pkgs/clan-cli/clan_cli/__init__.py b/pkgs/clan-cli/clan_cli/__init__.py index 7f2b9ba3..76b45591 100644 --- a/pkgs/clan-cli/clan_cli/__init__.py +++ b/pkgs/clan-cli/clan_cli/__init__.py @@ -152,6 +152,7 @@ For more detailed information, visit: https://docs.clan.lol/getting-started ), formatter_class=argparse.RawTextHelpFormatter, ) + flakes.register_parser(parser_flake) parser_config = subparsers.add_parser( diff --git a/pkgs/clan-cli/clan_cli/backups/create.py b/pkgs/clan-cli/clan_cli/backups/create.py index cb770c8e..7a2018c1 100644 --- a/pkgs/clan-cli/clan_cli/backups/create.py +++ b/pkgs/clan-cli/clan_cli/backups/create.py @@ -2,6 +2,7 @@ import argparse import json import logging +from ..completions import add_dynamic_completer, complete_machines from ..errors import ClanError from ..machines.machines import Machine @@ -40,8 +41,10 @@ def create_command(args: argparse.Namespace) -> None: def register_create_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machine", type=str, help="machine in the flake to create backups of" ) + add_dynamic_completer(machines_parser, complete_machines) + parser.add_argument("--provider", type=str, help="backup provider to use") parser.set_defaults(func=create_command) diff --git a/pkgs/clan-cli/clan_cli/backups/list.py b/pkgs/clan-cli/clan_cli/backups/list.py index 7ae15c27..9c466070 100644 --- a/pkgs/clan-cli/clan_cli/backups/list.py +++ b/pkgs/clan-cli/clan_cli/backups/list.py @@ -3,6 +3,7 @@ import json import subprocess from dataclasses import dataclass +from ..completions import add_dynamic_completer, complete_machines from ..errors import ClanError from ..machines.machines import Machine @@ -57,8 +58,9 @@ def list_command(args: argparse.Namespace) -> None: def register_list_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machine", type=str, help="machine in the flake to show backups of" ) + add_dynamic_completer(machines_parser, complete_machines) parser.add_argument("--provider", type=str, help="backup provider to filter by") parser.set_defaults(func=list_command) diff --git a/pkgs/clan-cli/clan_cli/completions.py b/pkgs/clan-cli/clan_cli/completions.py new file mode 100644 index 00000000..466c2926 --- /dev/null +++ b/pkgs/clan-cli/clan_cli/completions.py @@ -0,0 +1,81 @@ +import argparse +import json +import subprocess +import threading +from collections.abc import Callable, Iterable +from types import ModuleType +from typing import Any + +""" +This module provides dynamic completions. +The completions should feel fast. +We target a maximum of 1second on our average machine. +""" + + +argcomplete: ModuleType | None = None +try: + import argcomplete # type: ignore[no-redef] +except ImportError: + pass + + +def clan_dir(flake: str | None) -> str | None: + from .dirs import get_clan_flake_toplevel_or_env + + path_result = get_clan_flake_toplevel_or_env() + return str(path_result) if path_result is not None else None + + +def complete_machines( + prefix: str, parsed_args: argparse.Namespace, **kwargs: Any +) -> Iterable[str]: + """ + Provides completion functionality for machine names configured in the clan. + """ + machines: list[str] = [] + + def run_cmd() -> None: + try: + # In my tests this was consistently faster than: + # nix eval .#nixosConfigurations --apply builtins.attrNames + cmd = ["nix", "flake", "show", "--system", "no-eval", "--json"] + if (clan_dir_result := clan_dir(None)) is not None: + cmd.append(clan_dir_result) + result = subprocess.run( + cmd, + check=True, + capture_output=True, + text=True, + ) + + data = json.loads(result.stdout) + try: + machines.extend(data.get("nixosConfigurations").keys()) + except KeyError: + pass + except subprocess.CalledProcessError: + pass + + thread = threading.Thread(target=run_cmd) + thread.start() + thread.join(timeout=3) + + if thread.is_alive(): + return iter([]) + + machines_dict = {name: "machine" for name in machines} + return machines_dict + + +def add_dynamic_completer( + action: argparse.Action, + completer: Callable[..., Iterable[str]], +) -> None: + """ + Add a completion function to an argparse action, this will only be added, + if the argcomplete module is loaded. + """ + if argcomplete: + # mypy doesn't check this correctly, so we ignore it + action.completer = completer # type: ignore[attr-defined] diff --git a/pkgs/clan-cli/clan_cli/facts/check.py b/pkgs/clan-cli/clan_cli/facts/check.py index f15f195f..d879b2d7 100644 --- a/pkgs/clan-cli/clan_cli/facts/check.py +++ b/pkgs/clan-cli/clan_cli/facts/check.py @@ -2,6 +2,7 @@ import argparse import importlib import logging +from ..completions import add_dynamic_completer, complete_machines from ..machines.machines import Machine log = logging.getLogger(__name__) @@ -54,10 +55,12 @@ def check_command(args: argparse.Namespace) -> None: def register_check_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machine", help="The machine to check secrets for", ) + add_dynamic_completer(machines_parser, complete_machines) + parser.add_argument( "--service", help="the service to check", diff --git a/pkgs/clan-cli/clan_cli/facts/generate.py b/pkgs/clan-cli/clan_cli/facts/generate.py index dd033e1a..efcef4ba 100644 --- a/pkgs/clan-cli/clan_cli/facts/generate.py +++ b/pkgs/clan-cli/clan_cli/facts/generate.py @@ -9,6 +9,7 @@ from tempfile import TemporaryDirectory from clan_cli.cmd import run +from ..completions import add_dynamic_completer, complete_machines from ..errors import ClanError from ..git import commit_files from ..machines.inventory import get_all_machines, get_selected_machines @@ -216,13 +217,15 @@ def generate_command(args: argparse.Namespace) -> None: def register_generate_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machines", type=str, help="machine to generate facts for. if empty, generate facts for all machines", nargs="*", default=[], ) + add_dynamic_completer(machines_parser, complete_machines) + parser.add_argument( "--service", type=str, diff --git a/pkgs/clan-cli/clan_cli/facts/list.py b/pkgs/clan-cli/clan_cli/facts/list.py index 172f7d8a..0ec2fcdc 100644 --- a/pkgs/clan-cli/clan_cli/facts/list.py +++ b/pkgs/clan-cli/clan_cli/facts/list.py @@ -3,6 +3,7 @@ import importlib import json import logging +from ..completions import add_dynamic_completer, complete_machines from ..machines.machines import Machine log = logging.getLogger(__name__) @@ -37,8 +38,10 @@ def get_command(args: argparse.Namespace) -> None: def register_list_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machine", help="The machine to print facts for", ) + add_dynamic_completer(machines_parser, complete_machines) + parser.set_defaults(func=get_command) diff --git a/pkgs/clan-cli/clan_cli/facts/upload.py b/pkgs/clan-cli/clan_cli/facts/upload.py index 4e18a3bf..8d8cf30d 100644 --- a/pkgs/clan-cli/clan_cli/facts/upload.py +++ b/pkgs/clan-cli/clan_cli/facts/upload.py @@ -5,6 +5,7 @@ from pathlib import Path from tempfile import TemporaryDirectory from ..cmd import Log, run +from ..completions import add_dynamic_completer, complete_machines from ..machines.machines import Machine from ..nix import nix_shell @@ -46,8 +47,10 @@ def upload_command(args: argparse.Namespace) -> None: def register_upload_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machine", help="The machine to upload secrets to", ) + add_dynamic_completer(machines_parser, complete_machines) + parser.set_defaults(func=upload_command) diff --git a/pkgs/clan-cli/clan_cli/flash.py b/pkgs/clan-cli/clan_cli/flash.py index 91d4df49..85fcfd0c 100644 --- a/pkgs/clan-cli/clan_cli/flash.py +++ b/pkgs/clan-cli/clan_cli/flash.py @@ -12,6 +12,7 @@ from tempfile import TemporaryDirectory from typing import Any from .cmd import Log, run +from .completions import add_dynamic_completer, complete_machines from .errors import ClanError from .facts.secret_modules import SecretStoreBase from .machines.machines import Machine @@ -173,11 +174,13 @@ def flash_command(args: argparse.Namespace) -> None: def register_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machine", type=str, help="machine to install", ) + add_dynamic_completer(machines_parser, complete_machines) + parser.add_argument( "--disk", type=str, diff --git a/pkgs/clan-cli/clan_cli/machines/delete.py b/pkgs/clan-cli/clan_cli/machines/delete.py index 347ab65f..9291f957 100644 --- a/pkgs/clan-cli/clan_cli/machines/delete.py +++ b/pkgs/clan-cli/clan_cli/machines/delete.py @@ -1,6 +1,7 @@ import argparse import shutil +from ..completions import add_dynamic_completer, complete_machines from ..dirs import specific_machine_dir from ..errors import ClanError @@ -14,5 +15,7 @@ def delete_command(args: argparse.Namespace) -> None: def register_delete_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument("host", type=str) + machines_parser = parser.add_argument("host", type=str) + add_dynamic_completer(machines_parser, complete_machines) + parser.set_defaults(func=delete_command) diff --git a/pkgs/clan-cli/clan_cli/machines/install.py b/pkgs/clan-cli/clan_cli/machines/install.py index daf72759..183cc66e 100644 --- a/pkgs/clan-cli/clan_cli/machines/install.py +++ b/pkgs/clan-cli/clan_cli/machines/install.py @@ -8,6 +8,7 @@ from pathlib import Path from tempfile import TemporaryDirectory from ..cmd import Log, run +from ..completions import add_dynamic_completer, complete_machines from ..facts.generate import generate_facts from ..machines.machines import Machine from ..nix import nix_shell @@ -188,11 +189,14 @@ def register_install_parser(parser: argparse.ArgumentParser) -> None: help="do not ask for confirmation", default=False, ) - parser.add_argument( + + machines_parser = parser.add_argument( "machine", type=str, help="machine to install", ) + add_dynamic_completer(machines_parser, complete_machines) + parser.add_argument( "target_host", type=str, diff --git a/pkgs/clan-cli/clan_cli/machines/update.py b/pkgs/clan-cli/clan_cli/machines/update.py index 1da38904..3a774ee9 100644 --- a/pkgs/clan-cli/clan_cli/machines/update.py +++ b/pkgs/clan-cli/clan_cli/machines/update.py @@ -6,6 +6,7 @@ import shlex import subprocess import sys +from ..completions import add_dynamic_completer, complete_machines from ..errors import ClanError from ..facts.generate import generate_facts from ..facts.upload import upload_secrets @@ -180,7 +181,7 @@ def update(args: argparse.Namespace) -> None: def register_update_parser(parser: argparse.ArgumentParser) -> None: - parser.add_argument( + machines_parser = parser.add_argument( "machines", type=str, nargs="*", @@ -188,6 +189,9 @@ def register_update_parser(parser: argparse.ArgumentParser) -> None: metavar="MACHINE", help="machine to update. If no machine is specified, all machines will be updated.", ) + + add_dynamic_completer(machines_parser, complete_machines) + parser.add_argument( "--target-host", type=str,