From e2ca4d7fd44d309a25b076a1701f194b9dcca453 Mon Sep 17 00:00:00 2001 From: Bernardino Date: Fri, 16 Jan 2026 16:58:23 +0000 Subject: [PATCH] Add tab autocompletion to the cli --- README.md | 52 +++++++ pyproject.toml | 1 + src/hive_cli/completers.py | 183 ++++++++++++++++++++++++ src/hive_cli/main.py | 23 +-- tests/test_completers.py | 280 +++++++++++++++++++++++++++++++++++++ 5 files changed, 529 insertions(+), 10 deletions(-) create mode 100644 src/hive_cli/completers.py create mode 100644 tests/test_completers.py diff --git a/README.md b/README.md index 85fb685..1760676 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,58 @@ pip install hiverge-cli source start.sh ``` +## Shell Completion + +Hive CLI supports shell tab completion for commands, experiment names, sandbox names, and file paths. After installing the CLI, enable completion for your shell: + +### Bash + +Add the following to your `~/.bashrc`: + +```bash +eval "$(register-python-argcomplete hive)" +``` + +Then reload your shell configuration: + +```bash +source ~/.bashrc +``` + +### Zsh + +Add the following to your `~/.zshrc`: + +```bash +autoload -U bashcompinit +bashcompinit +eval "$(register-python-argcomplete hive)" +``` + +Then reload your shell configuration: + +```bash +source ~/.zshrc +``` + +### Verify + +Test that completion is working: + +```bash +hive delete exp # Should list available experiments +hive log # Should list available sandbox pods +``` + +### Features + +- Command and subcommand completion for all hive commands +- Dynamic completion of experiment names for `delete` and `show sandboxes --experiment` +- Dynamic completion of sandbox pod names for `log` command +- File path completion for `-f/--config` flags + +**Note:** Completion requires access to your Kubernetes cluster to fetch experiment and sandbox names. If the cluster is unavailable, completion will still work for static options but won't show dynamic resources. + ## How to run **Note**: Hive-CLI reads the configuration from a yaml file, by default it will look for the `~/.hive/hive.yaml`. You can also specify a different configuration file using the `-f` option. Refer to the [hive.yaml](./hive.yaml) for examples. diff --git a/pyproject.toml b/pyproject.toml index a2e9b0a..6ff8778 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "rich>=12.5.1", "kubernetes>=33.1.0", "portforward>=0.3.0", + "argcomplete>=3.0.0", ] [dependency-groups] diff --git a/src/hive_cli/completers.py b/src/hive_cli/completers.py new file mode 100644 index 0000000..0be7bc4 --- /dev/null +++ b/src/hive_cli/completers.py @@ -0,0 +1,183 @@ +""" +Argcomplete completer functions for the Hive CLI. + +These functions provide dynamic tab completion for commands, fetching +resource names from Kubernetes and providing file path completion. +""" + +import os +import signal +from functools import wraps + +from argcomplete.completers import FilesCompleter + + +def safe_completer(func): + """ + Decorator to wrap completer functions with error handling. + + Ensures that completers never raise exceptions that would break the CLI. + All errors are silently caught and an empty list is returned. + """ + @wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except Exception: + # Silent failure - return empty list on any error + return [] + return wrapper + + +class TimeoutError(Exception): + """Raised when a completion operation times out.""" + pass + + +def timeout_handler(signum, frame): + """Signal handler for timeout.""" + raise TimeoutError("Operation timed out") + + +@safe_completer +def experiment_completer(prefix, parsed_args, **kwargs): + """ + Complete experiment names by fetching from Kubernetes. + + Used for: + - hive delete experiment + - hive show sandboxes --experiment + + Args: + prefix: The current prefix being completed + parsed_args: Parsed arguments from argparse + **kwargs: Additional arguments from argcomplete + + Returns: + List of experiment names matching the prefix + """ + # Set up 2-second timeout + signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(2) + + try: + # Import here to avoid loading K8s client at module import time + from hive_cli.config import load_config + from hive_cli.platform.k8s import K8sPlatform + + # Load config - use default if not specified + config_path = getattr(parsed_args, "config", None) + if not config_path: + config_path = os.path.expandvars("$HOME/.hive/hive.yaml") + + config = load_config(config_path) + + # Create platform and fetch experiments + platform = K8sPlatform(None, config.token_path) + resp = platform.client.list_namespaced_custom_object( + group="core.hiverge.ai", + version="v1alpha1", + namespace="default", + plural="experiments", + ) + + # Extract experiment names + experiments = [item["metadata"]["name"] for item in resp.get("items", [])] + + # Filter by prefix if provided + if prefix: + experiments = [exp for exp in experiments if exp.startswith(prefix)] + + return experiments + + except TimeoutError: + # Timeout - return empty list + return [] + finally: + # Cancel alarm + signal.alarm(0) + + +@safe_completer +def sandbox_completer(prefix, parsed_args, **kwargs): + """ + Complete sandbox pod names by fetching from Kubernetes. + + Used for: + - hive log + + If --experiment flag is provided, filters sandboxes by experiment label. + + Args: + prefix: The current prefix being completed + parsed_args: Parsed arguments from argparse + **kwargs: Additional arguments from argcomplete + + Returns: + List of sandbox pod names matching the prefix + """ + # Set up 2-second timeout + signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(2) + + try: + # Import here to avoid loading K8s client at module import time + from hive_cli.config import load_config + from hive_cli.platform.k8s import K8sPlatform + + # Load config - use default if not specified + config_path = getattr(parsed_args, "config", None) + if not config_path: + config_path = os.path.expandvars("$HOME/.hive/hive.yaml") + + config = load_config(config_path) + + # Create platform and fetch sandboxes + platform = K8sPlatform(None, config.token_path) + + # Build label selector + label_selector = "app=hive-sandbox" + + # Filter by experiment if provided + experiment = getattr(parsed_args, "experiment", None) + if experiment: + label_selector += f",hiverge.ai/experiment-name={experiment}" + + # Fetch pods + pods = platform.core_client.list_namespaced_pod( + namespace="default", + label_selector=label_selector + ) + + # Extract sandbox names + sandboxes = [pod.metadata.name for pod in pods.items] + + # Filter by prefix if provided + if prefix: + sandboxes = [sb for sb in sandboxes if sb.startswith(prefix)] + + return sandboxes + + except TimeoutError: + # Timeout - return empty list + return [] + finally: + # Cancel alarm + signal.alarm(0) + + +def config_file_completer(prefix, parsed_args, **kwargs): + """ + Complete config file paths for -f/--config flags. + + Delegates to argcomplete's built-in FilesCompleter for file path completion. + + Args: + prefix: The current prefix being completed + parsed_args: Parsed arguments from argparse + **kwargs: Additional arguments from argcomplete + + Returns: + List of file paths matching the prefix + """ + return FilesCompleter()(prefix, **kwargs) diff --git a/src/hive_cli/main.py b/src/hive_cli/main.py index 0b81695..cae16f6 100644 --- a/src/hive_cli/main.py +++ b/src/hive_cli/main.py @@ -3,10 +3,12 @@ import subprocess from importlib.metadata import PackageNotFoundError, version +import argcomplete import portforward from rich.console import Console from rich.text import Text +from hive_cli.completers import config_file_completer, experiment_completer, sandbox_completer from hive_cli.config import HiveConfig, load_config from hive_cli.platform.k8s import K8sPlatform from hive_cli.utils import event @@ -164,7 +166,7 @@ def main(): "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, default to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_create_exp.set_defaults(func=create_experiment_cli) # TODO: @@ -190,13 +192,13 @@ def main(): parser_delete_exp = delete_subparsers.add_parser( "experiment", aliases=["exp"], help="Delete an experiment" ) - parser_delete_exp.add_argument("name", help="Name of the experiment") + parser_delete_exp.add_argument("name", help="Name of the experiment").completer = experiment_completer parser_delete_exp.add_argument( "-f", "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, default to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_delete_exp.set_defaults(func=delete_experiment_cli) # show command @@ -212,7 +214,7 @@ def main(): "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, default to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_show_exp.set_defaults(func=show_experiment_cli) ## show sandboxes @@ -224,12 +226,12 @@ def main(): "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, default to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_show_sandbox.add_argument( "-exp", "--experiment", help="Name of the experiment running sandboxes", - ) + ).completer = experiment_completer parser_show_sandbox.set_defaults(func=show_sandbox_cli) # edit command @@ -243,7 +245,7 @@ def main(): "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, defaults to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_edit_config.set_defaults(func=edit_cli) # dashboard command @@ -259,7 +261,7 @@ def main(): "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, default to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_dashboard.set_defaults(func=show_dashboard_cli) # version command @@ -268,13 +270,13 @@ def main(): # log command parser_log = subparsers.add_parser("log", help="Show Sandbox logs") - parser_log.add_argument("sandbox", help="Name of the sandbox to fetch logs for") + parser_log.add_argument("sandbox", help="Name of the sandbox to fetch logs for").completer = sandbox_completer parser_log.add_argument( "-f", "--config", default=os.path.expandvars("$HOME/.hive/hive.yaml"), help="Path to the config file, default to ~/.hive/hive.yaml", - ) + ).completer = config_file_completer parser_log.add_argument( "-t", "--tail", @@ -284,6 +286,7 @@ def main(): ) parser_log.set_defaults(func=display_sandbox_logs_cli) + argcomplete.autocomplete(parser) args = parser.parse_args() if hasattr(args, "func"): args.func(args) diff --git a/tests/test_completers.py b/tests/test_completers.py new file mode 100644 index 0000000..44eabcb --- /dev/null +++ b/tests/test_completers.py @@ -0,0 +1,280 @@ +""" +Unit tests for argcomplete completer functions. +""" + +from argparse import Namespace +from unittest.mock import MagicMock, patch + +from hive_cli.completers import ( + config_file_completer, + experiment_completer, + sandbox_completer, +) + + +class TestExperimentCompleter: + """Tests for experiment_completer function.""" + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_experiment_completer_success(self, mock_platform_class, mock_load_config): + """Test experiment_completer successfully fetches and returns experiment names.""" + # Setup mock config + mock_config = MagicMock() + mock_config.token_path = "/path/to/token" + mock_load_config.return_value = mock_config + + # Setup mock platform + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + # Mock K8s API response + mock_platform.client.list_namespaced_custom_object.return_value = { + "items": [ + {"metadata": {"name": "exp-1"}}, + {"metadata": {"name": "exp-2"}}, + {"metadata": {"name": "exp-3"}}, + ] + } + + # Create parsed args + parsed_args = Namespace(config=None) + + # Call completer + result = experiment_completer("", parsed_args) + + # Verify results + assert result == ["exp-1", "exp-2", "exp-3"] + mock_platform.client.list_namespaced_custom_object.assert_called_once_with( + group="core.hiverge.ai", + version="v1alpha1", + namespace="default", + plural="experiments", + ) + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_experiment_completer_with_prefix(self, mock_platform_class, mock_load_config): + """Test experiment_completer filters by prefix.""" + # Setup mocks + mock_config = MagicMock() + mock_config.token_path = "/path/to/token" + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + mock_platform.client.list_namespaced_custom_object.return_value = { + "items": [ + {"metadata": {"name": "exp-1"}}, + {"metadata": {"name": "exp-2"}}, + {"metadata": {"name": "test-1"}}, + ] + } + + parsed_args = Namespace(config=None) + + # Call completer with prefix + result = experiment_completer("exp", parsed_args) + + # Should only return experiments starting with "exp" + assert result == ["exp-1", "exp-2"] + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_experiment_completer_k8s_error(self, mock_platform_class, mock_load_config): + """Test experiment_completer returns empty list on K8s API error.""" + # Setup mocks + mock_config = MagicMock() + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + # Simulate K8s API error + mock_platform.client.list_namespaced_custom_object.side_effect = Exception("K8s error") + + parsed_args = Namespace(config=None) + + # Call completer + result = experiment_completer("", parsed_args) + + # Should return empty list on error + assert result == [] + + @patch("hive_cli.config.load_config") + def test_experiment_completer_config_error(self, mock_load_config): + """Test experiment_completer returns empty list when config cannot be loaded.""" + # Simulate config loading error + mock_load_config.side_effect = Exception("Config error") + + parsed_args = Namespace(config=None) + + # Call completer + result = experiment_completer("", parsed_args) + + # Should return empty list on error + assert result == [] + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_experiment_completer_empty_response(self, mock_platform_class, mock_load_config): + """Test experiment_completer handles empty K8s response.""" + # Setup mocks + mock_config = MagicMock() + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + # Mock empty response + mock_platform.client.list_namespaced_custom_object.return_value = {"items": []} + + parsed_args = Namespace(config=None) + + # Call completer + result = experiment_completer("", parsed_args) + + # Should return empty list + assert result == [] + + +class TestSandboxCompleter: + """Tests for sandbox_completer function.""" + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_sandbox_completer_success(self, mock_platform_class, mock_load_config): + """Test sandbox_completer successfully fetches and returns sandbox names.""" + # Setup mocks + mock_config = MagicMock() + mock_config.token_path = "/path/to/token" + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + # Mock pods response + mock_pod1 = MagicMock() + mock_pod1.metadata.name = "sandbox-1" + mock_pod2 = MagicMock() + mock_pod2.metadata.name = "sandbox-2" + + mock_pods = MagicMock() + mock_pods.items = [mock_pod1, mock_pod2] + mock_platform.core_client.list_namespaced_pod.return_value = mock_pods + + parsed_args = Namespace(config=None, experiment=None) + + # Call completer + result = sandbox_completer("", parsed_args) + + # Verify results + assert result == ["sandbox-1", "sandbox-2"] + mock_platform.core_client.list_namespaced_pod.assert_called_once_with( + namespace="default", + label_selector="app=hive-sandbox" + ) + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_sandbox_completer_with_experiment_filter(self, mock_platform_class, mock_load_config): + """Test sandbox_completer filters by experiment when provided.""" + # Setup mocks + mock_config = MagicMock() + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + mock_pod = MagicMock() + mock_pod.metadata.name = "sandbox-1" + + mock_pods = MagicMock() + mock_pods.items = [mock_pod] + mock_platform.core_client.list_namespaced_pod.return_value = mock_pods + + parsed_args = Namespace(config=None, experiment="exp-1") + + # Call completer + result = sandbox_completer("", parsed_args) + + # Verify experiment filter was applied + assert result == ["sandbox-1"] + mock_platform.core_client.list_namespaced_pod.assert_called_once_with( + namespace="default", + label_selector="app=hive-sandbox,hiverge.ai/experiment-name=exp-1" + ) + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_sandbox_completer_with_prefix(self, mock_platform_class, mock_load_config): + """Test sandbox_completer filters by prefix.""" + # Setup mocks + mock_config = MagicMock() + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + mock_pod1 = MagicMock() + mock_pod1.metadata.name = "sandbox-1" + mock_pod2 = MagicMock() + mock_pod2.metadata.name = "sandbox-2" + mock_pod3 = MagicMock() + mock_pod3.metadata.name = "test-pod" + + mock_pods = MagicMock() + mock_pods.items = [mock_pod1, mock_pod2, mock_pod3] + mock_platform.core_client.list_namespaced_pod.return_value = mock_pods + + parsed_args = Namespace(config=None, experiment=None) + + # Call completer with prefix + result = sandbox_completer("sandbox", parsed_args) + + # Should only return sandboxes starting with "sandbox" + assert result == ["sandbox-1", "sandbox-2"] + + @patch("hive_cli.config.load_config") + @patch("hive_cli.platform.k8s.K8sPlatform") + def test_sandbox_completer_k8s_error(self, mock_platform_class, mock_load_config): + """Test sandbox_completer returns empty list on K8s API error.""" + # Setup mocks + mock_config = MagicMock() + mock_load_config.return_value = mock_config + + mock_platform = MagicMock() + mock_platform_class.return_value = mock_platform + + # Simulate K8s API error + mock_platform.core_client.list_namespaced_pod.side_effect = Exception("K8s error") + + parsed_args = Namespace(config=None, experiment=None) + + # Call completer + result = sandbox_completer("", parsed_args) + + # Should return empty list on error + assert result == [] + + +class TestConfigFileCompleter: + """Tests for config_file_completer function.""" + + @patch("hive_cli.completers.FilesCompleter") + def test_config_file_completer(self, mock_files_completer): + """Test config_file_completer delegates to FilesCompleter.""" + # Setup mock + mock_completer_instance = MagicMock() + mock_completer_instance.return_value = ["/path/to/file1.yaml", "/path/to/file2.yaml"] + mock_files_completer.return_value = mock_completer_instance + + parsed_args = Namespace() + + # Call completer + result = config_file_completer("~/.hive/", parsed_args) + + # Verify FilesCompleter was used + mock_files_completer.assert_called_once() + assert result == ["/path/to/file1.yaml", "/path/to/file2.yaml"]