From 1abcc3e71b032da92b7f5b35a65a9e3eccd88256 Mon Sep 17 00:00:00 2001
From: Zhe Yu <zcabzyu@ucl.ac.uk>
Date: Tue, 24 Jun 2025 18:02:31 +0800
Subject: [PATCH] fix(cli): Fix shtab generation for zsh and add tests for it.
 Fix #211

---
 src/vectorcode/cli_utils.py |  2 +-
 tests/test_cli_utils.py     | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/vectorcode/cli_utils.py b/src/vectorcode/cli_utils.py
index 5fa23b3c..38fd74ef 100644
--- a/src/vectorcode/cli_utils.py
+++ b/src/vectorcode/cli_utils.py
@@ -358,7 +358,7 @@ def get_cli_parser():
     )
     prompts_parser.add_argument(
         "prompt_categories",
-        choices=PromptCategory,
+        choices=[str(i) for i in PromptCategory],
         type=PromptCategory,
         nargs="*",
         help="The subcommand(s) to get the prompts for. When not provided, VectorCode will print the prompts for `query`.",
diff --git a/tests/test_cli_utils.py b/tests/test_cli_utils.py
index 6f6f09f7..3252683f 100644
--- a/tests/test_cli_utils.py
+++ b/tests/test_cli_utils.py
@@ -1,4 +1,6 @@
 import os
+import subprocess
+import sys
 import tempfile
 from typing import Any, Dict
 from unittest.mock import patch
@@ -539,3 +541,15 @@ def test_cleanup_path():
         "~", "test_path"
     )
     assert cleanup_path("/etc/dir") == "/etc/dir"
+
+
+def test_shtab():
+    for shell in ("bash", "zsh", "tcsh"):
+        assert (
+            subprocess.Popen(
+                [sys.executable, "-m", "vectorcode.main", "-s", shell],
+                stderr=subprocess.PIPE,
+            )
+            .stderr.read()
+            .decode()
+        ) == ""