Files
flow/commands/completion.py
2026-02-12 09:42:59 +02:00

526 lines
15 KiB
Python

"""flow completion — shell completion support (dynamic zsh)."""
import argparse
import json
import shutil
import subprocess
from pathlib import Path
from typing import List, Optional, Sequence, Set
from flow.commands.enter import HOST_TEMPLATES
from flow.core.config import load_config, load_manifest
from flow.core.paths import DOTFILES_DIR, INSTALLED_STATE
ZSH_RC_START = "# >>> flow completion >>>"
ZSH_RC_END = "# <<< flow completion <<<"
TOP_LEVEL_COMMANDS = [
"enter",
"dev",
"dotfiles",
"dot",
"bootstrap",
"setup",
"provision",
"package",
"pkg",
"sync",
"completion",
]
def register(subparsers):
p = subparsers.add_parser("completion", help="Shell completion helpers")
sub = p.add_subparsers(dest="completion_command")
zsh = sub.add_parser("zsh", help="Print zsh completion script")
zsh.set_defaults(handler=run_zsh_script)
install = sub.add_parser("install-zsh", help="Install zsh completion script")
install.add_argument(
"--dir",
default="~/.zsh/completions",
help="Directory where _flow completion file is written",
)
install.add_argument(
"--rc",
default="~/.zshrc",
help="Shell rc file to update with fpath/compinit snippet",
)
install.add_argument(
"--no-rc",
action="store_true",
help="Do not modify rc file; only write completion script",
)
install.set_defaults(handler=run_install_zsh)
hidden = sub.add_parser("_zsh_complete", help=argparse.SUPPRESS)
hidden.add_argument("--cword", type=int, required=True, help=argparse.SUPPRESS)
hidden.add_argument("words", nargs="*", help=argparse.SUPPRESS)
hidden.set_defaults(handler=run_zsh_complete)
p.set_defaults(handler=lambda _ctx, args: p.print_help())
def _canonical_command(command: str) -> str:
alias_map = {
"dot": "dotfiles",
"setup": "bootstrap",
"provision": "bootstrap",
"pkg": "package",
}
return alias_map.get(command, command)
def _safe_config():
try:
return load_config()
except Exception:
return None
def _safe_manifest():
try:
return load_manifest()
except Exception:
return {}
def _list_targets() -> List[str]:
cfg = _safe_config()
if cfg is None:
return []
return sorted({f"{t.namespace}@{t.platform}" for t in cfg.targets})
def _list_namespaces() -> List[str]:
cfg = _safe_config()
if cfg is None:
return []
return sorted({t.namespace for t in cfg.targets})
def _list_platforms() -> List[str]:
cfg = _safe_config()
config_platforms: Set[str] = set()
if cfg is not None:
config_platforms = {t.platform for t in cfg.targets}
return sorted(set(HOST_TEMPLATES.keys()) | config_platforms)
def _list_bootstrap_profiles() -> List[str]:
manifest = _safe_manifest()
return sorted(manifest.get("profiles", {}).keys())
def _list_manifest_packages() -> List[str]:
manifest = _safe_manifest()
return sorted(manifest.get("binaries", {}).keys())
def _list_installed_packages() -> List[str]:
if not INSTALLED_STATE.exists():
return []
try:
with open(INSTALLED_STATE) as f:
state = json.load(f)
except Exception:
return []
if not isinstance(state, dict):
return []
return sorted(state.keys())
def _list_dotfiles_profiles() -> List[str]:
profiles_dir = DOTFILES_DIR / "profiles"
if not profiles_dir.is_dir():
return []
return sorted([p.name for p in profiles_dir.iterdir() if p.is_dir() and not p.name.startswith(".")])
def _list_dotfiles_packages(profile: Optional[str] = None) -> List[str]:
package_names: Set[str] = set()
common = DOTFILES_DIR / "common"
if common.is_dir():
for pkg in common.iterdir():
if pkg.is_dir() and not pkg.name.startswith("."):
package_names.add(pkg.name)
if profile:
profile_dir = DOTFILES_DIR / "profiles" / profile
if profile_dir.is_dir():
for pkg in profile_dir.iterdir():
if pkg.is_dir() and not pkg.name.startswith("."):
package_names.add(pkg.name)
else:
profiles_dir = DOTFILES_DIR / "profiles"
if profiles_dir.is_dir():
for profile_dir in profiles_dir.iterdir():
if not profile_dir.is_dir():
continue
for pkg in profile_dir.iterdir():
if pkg.is_dir() and not pkg.name.startswith("."):
package_names.add(pkg.name)
return sorted(package_names)
def _list_container_names() -> List[str]:
runtime = None
for rt in ("docker", "podman"):
if shutil.which(rt):
runtime = rt
break
if not runtime:
return []
try:
result = subprocess.run(
[
runtime,
"ps",
"-a",
"--filter",
"label=dev=true",
"--format",
'{{.Label "dev.name"}}',
],
capture_output=True,
text=True,
timeout=1,
)
except Exception:
return []
if result.returncode != 0:
return []
names = []
for line in result.stdout.splitlines():
line = line.strip()
if line:
names.append(line)
return sorted(set(names))
def _split_words(words: Sequence[str], cword: int):
tokens = list(words)
index = max(0, cword - 1)
if tokens:
tokens = tokens[1:]
index = max(0, cword - 2)
if index > len(tokens):
index = len(tokens)
current = tokens[index] if index < len(tokens) else ""
before = tokens[:index]
return before, current
def _filter(candidates: Sequence[str], prefix: str) -> List[str]:
unique = sorted(set(candidates))
if not prefix:
return unique
return [c for c in unique if c.startswith(prefix)]
def _profile_from_before(before: Sequence[str]) -> Optional[str]:
for i, token in enumerate(before):
if token == "--profile" and i + 1 < len(before):
return before[i + 1]
return None
def _complete_dev(before: Sequence[str], current: str) -> List[str]:
if len(before) <= 1:
return _filter(["create", "exec", "connect", "list", "stop", "remove", "rm", "respawn"], current)
sub = "remove" if before[1] == "rm" else before[1]
if sub in {"remove", "stop", "connect", "exec", "respawn"}:
options = {
"remove": ["-f", "--force", "-h", "--help"],
"stop": ["--kill", "-h", "--help"],
"exec": ["-h", "--help"],
"connect": ["-h", "--help"],
"respawn": ["-h", "--help"],
}[sub]
if current.startswith("-"):
return _filter(options, current)
non_opt = [t for t in before[2:] if not t.startswith("-")]
if len(non_opt) == 0:
return _filter(_list_container_names(), current)
return []
if sub == "create":
options = ["-i", "--image", "-p", "--project", "-h", "--help"]
if before and before[-1] in ("-i", "--image"):
return _filter(["tm0/node", "docker/python", "docker/alpine"], current)
if current.startswith("-"):
return _filter(options, current)
return []
if sub == "list":
return []
return []
def _complete_dotfiles(before: Sequence[str], current: str) -> List[str]:
if len(before) <= 1:
return _filter(
["init", "link", "unlink", "status", "sync", "relink", "clean", "edit"],
current,
)
sub = before[1]
if sub == "init":
return _filter(["--repo", "-h", "--help"], current) if current.startswith("-") else []
if sub in {"link", "relink"}:
if before and before[-1] == "--profile":
return _filter(_list_dotfiles_profiles(), current)
if current.startswith("-"):
return _filter(["--profile", "--copy", "--force", "--dry-run", "-h", "--help"], current)
profile = _profile_from_before(before)
return _filter(_list_dotfiles_packages(profile), current)
if sub == "unlink":
if current.startswith("-"):
return _filter(["-h", "--help"], current)
return _filter(_list_dotfiles_packages(), current)
if sub == "edit":
if current.startswith("-"):
return _filter(["--no-commit", "-h", "--help"], current)
non_opt = [t for t in before[2:] if not t.startswith("-")]
if len(non_opt) == 0:
return _filter(_list_dotfiles_packages(), current)
return []
if sub == "clean":
return _filter(["--dry-run", "-h", "--help"], current) if current.startswith("-") else []
return []
def _complete_bootstrap(before: Sequence[str], current: str) -> List[str]:
if len(before) <= 1:
return _filter(["run", "list", "show"], current)
sub = before[1]
if sub == "run":
if before and before[-1] == "--profile":
return _filter(_list_bootstrap_profiles(), current)
if current.startswith("-"):
return _filter(["--profile", "--dry-run", "--var", "-h", "--help"], current)
return []
if sub == "show":
if current.startswith("-"):
return _filter(["-h", "--help"], current)
non_opt = [t for t in before[2:] if not t.startswith("-")]
if len(non_opt) == 0:
return _filter(_list_bootstrap_profiles(), current)
return []
return []
def _complete_package(before: Sequence[str], current: str) -> List[str]:
if len(before) <= 1:
return _filter(["install", "list", "remove"], current)
sub = before[1]
if sub == "install":
if current.startswith("-"):
return _filter(["--dry-run", "-h", "--help"], current)
return _filter(_list_manifest_packages(), current)
if sub == "remove":
if current.startswith("-"):
return _filter(["-h", "--help"], current)
return _filter(_list_installed_packages(), current)
if sub == "list":
if current.startswith("-"):
return _filter(["--all", "-h", "--help"], current)
return []
return []
def _complete_sync(before: Sequence[str], current: str) -> List[str]:
if len(before) <= 1:
return _filter(["check", "fetch", "summary"], current)
sub = before[1]
if sub == "check" and current.startswith("-"):
return _filter(["--fetch", "--no-fetch", "-h", "--help"], current)
if current.startswith("-"):
return _filter(["-h", "--help"], current)
return []
def complete(words: Sequence[str], cword: int) -> List[str]:
before, current = _split_words(words, cword)
if not before:
return _filter(TOP_LEVEL_COMMANDS + ["-h", "--help", "-v", "--version"], current)
command = _canonical_command(before[0])
if command == "enter":
if before and before[-1] in ("-p", "--platform"):
return _filter(_list_platforms(), current)
if before and before[-1] in ("-n", "--namespace"):
return _filter(_list_namespaces(), current)
if current.startswith("-"):
return _filter(
["-u", "--user", "-n", "--namespace", "-p", "--platform", "-s", "--session", "--no-tmux", "-d", "--dry-run", "-h", "--help"],
current,
)
return _filter(_list_targets(), current)
if command == "dev":
return _complete_dev(before, current)
if command == "dotfiles":
return _complete_dotfiles(before, current)
if command == "bootstrap":
return _complete_bootstrap(before, current)
if command == "package":
return _complete_package(before, current)
if command == "sync":
return _complete_sync(before, current)
if command == "completion":
if len(before) <= 1:
return _filter(["zsh", "install-zsh"], current)
sub = before[1]
if sub == "install-zsh" and current.startswith("-"):
return _filter(["--dir", "--rc", "--no-rc", "-h", "--help"], current)
return []
return []
def run_zsh_complete(_ctx, args):
candidates = complete(args.words, args.cword)
for item in candidates:
print(item)
def _zsh_script_text() -> str:
return r'''#compdef flow
_flow() {
local -a suggestions
suggestions=("${(@f)$(flow completion _zsh_complete --cword "$CURRENT" -- "${words[@]}" 2>/dev/null)}")
if (( ${#suggestions[@]} > 0 )); then
compadd -Q -- "${suggestions[@]}"
return 0
fi
if [[ "$words[CURRENT]" == */* || "$words[CURRENT]" == ./* || "$words[CURRENT]" == ~* ]]; then
_files
return 0
fi
return 1
}
compdef _flow flow
'''
def _zsh_dir_for_rc(path: Path) -> str:
home = Path.home().resolve()
resolved = path.expanduser().resolve()
try:
rel = resolved.relative_to(home)
return f"~/{rel}" if str(rel) != "." else "~"
except ValueError:
return str(resolved)
def _zsh_rc_snippet(completions_dir: Path) -> str:
dir_expr = _zsh_dir_for_rc(completions_dir)
return (
f"{ZSH_RC_START}\n"
f"fpath=({dir_expr} $fpath)\n"
"autoload -Uz compinit && compinit\n"
f"{ZSH_RC_END}\n"
)
def _ensure_rc_snippet(rc_path: Path, completions_dir: Path) -> bool:
snippet = _zsh_rc_snippet(completions_dir)
if rc_path.exists():
content = rc_path.read_text()
else:
content = ""
if ZSH_RC_START in content and ZSH_RC_END in content:
start = content.find(ZSH_RC_START)
end = content.find(ZSH_RC_END, start)
if end >= 0:
end += len(ZSH_RC_END)
updated = content[:start] + snippet.rstrip("\n") + content[end:]
if updated == content:
return False
rc_path.parent.mkdir(parents=True, exist_ok=True)
rc_path.write_text(updated)
return True
sep = "" if content.endswith("\n") or content == "" else "\n"
rc_path.parent.mkdir(parents=True, exist_ok=True)
rc_path.write_text(content + sep + snippet)
return True
def run_install_zsh(_ctx, args):
completions_dir = Path(args.dir).expanduser()
completions_dir.mkdir(parents=True, exist_ok=True)
completion_file = completions_dir / "_flow"
completion_file.write_text(_zsh_script_text())
print(f"Installed completion script: {completion_file}")
if args.no_rc:
print("Skipped rc file update (--no-rc)")
return
rc_path = Path(args.rc).expanduser()
changed = _ensure_rc_snippet(rc_path, completions_dir)
if changed:
print(f"Updated shell rc: {rc_path}")
else:
print(f"Shell rc already configured: {rc_path}")
print("Restart shell or run: autoload -Uz compinit && compinit")
def run_zsh_script(_ctx, _args):
print(_zsh_script_text())