implement autocompletion for nodes

comfy_cli/cmdline.py

@@ -252,7 +252,13 @@ def install(
 
 @app.command(help="Update ComfyUI Environment [all|comfy]")
 @tracking.track_command()
-def update(target: str = typer.Argument("comfy", help="[all|comfy]")):
+def update(
+    target: str = typer.Argument(
+        "comfy",
+        help="[all|comfy]",
+        autocompletion=utils.create_choice_completer(["all", "comfy"]),
+    )
+):
     if target not in ["all", "comfy"]:
         typer.echo(
             f"Invalid target: {target}. Allowed targets are 'all', 'comfy'.",
@@ -277,6 +283,8 @@ def update(target: str = typer.Argument("comfy", help="[all|comfy]")):
             check=True,
         )
 
+    custom_nodes.command.update_node_id_cache()
+
 
 # @app.command(help="Run workflow file")
 # @tracking.track_command()

comfy_cli/command/custom_nodes/command.py

@@ -9,7 +9,7 @@
 from rich import print
 from typing_extensions import List, Annotated
 
-from comfy_cli import ui, logging, tracking
+from comfy_cli import ui, logging, tracking, utils
 from comfy_cli.config_manager import ConfigManager
 from comfy_cli.file_utils import (
     download_file,
@@ -168,8 +168,9 @@ def clear(path: str):
     execute_cm_cli(["clear", path])
 
 
-def show_completer(incomplete: str) -> list[str]:
-    valid_choices = [
+# completers
+show_completer = utils.create_choice_completer(
+    [
         "installed",
         "enabled",
         "not-installed",
@@ -178,17 +179,33 @@ def show_completer(incomplete: str) -> list[str]:
         "snapshot",
         "snapshot-list",
     ]
-    return [choice for choice in valid_choices if choice.startswith(incomplete)]
+)
+
+
+mode_completer = utils.create_choice_completer(["remote", "local", "cache"])
+
+
+channel_completer = utils.create_choice_completer(
+    ["default", "recent", "dev", "forked", "tutorial", "legacy"]
+)
 
 
-def mode_completer(incomplete: str) -> list[str]:
-    modes = ["remote", "local", "cache"]
-    return [mode for mode in modes if mode.startswith(incomplete)]
+def node_completer(incomplete: str) -> list[str]:
+    try:
+        config_manager = ConfigManager()
+        tmp_path = os.path.join(
+            config_manager.get_config_path(), "tmp", "node-cache.list"
+        )
 
+        with open(tmp_path, "r", encoding="UTF-8", errors="ignore") as cache_file:
+            return [
+                node_id
+                for node_id in cache_file.readlines()
+                if node_id.startswith(incomplete)
+            ]
 
-def channel_completer(incomplete: str) -> list[str]:
-    modes = ["default", "recent", "dev", "forked", "tutorial", "legacy"]
-    return [mode for mode in modes if mode.startswith(incomplete)]
+    except Exception:
+        return []
 
 
 @app.command(help="Show node list")
@@ -285,7 +302,9 @@ def simple_show(
 @app.command(help="Install custom nodes")
 @tracking.track_command("node")
 def install(
-    args: List[str] = typer.Argument(..., help="install custom nodes"),
+    args: List[str] = typer.Argument(
+        ..., help="install custom nodes", autocompletion=node_completer
+    ),
     channel: Annotated[
         str,
         typer.Option(
@@ -318,7 +337,9 @@ def install(
 @app.command(help="Reinstall custom nodes")
 @tracking.track_command("node")
 def reinstall(
-    args: List[str] = typer.Argument(..., help="reinstall custom nodes"),
+    args: List[str] = typer.Argument(
+        ..., help="reinstall custom nodes", autocompletion=node_completer
+    ),
     channel: Annotated[
         str,
         typer.Option(
@@ -351,7 +372,9 @@ def reinstall(
 @app.command(help="Uninstall custom nodes")
 @tracking.track_command("node")
 def uninstall(
-    args: List[str] = typer.Argument(..., help="uninstall custom nodes"),
+    args: List[str] = typer.Argument(
+        ..., help="uninstall custom nodes", autocompletion=node_completer
+    ),
     channel: Annotated[
         str,
         typer.Option(
@@ -381,13 +404,36 @@ def uninstall(
     execute_cm_cli(["uninstall"] + args, channel, mode)
 
 
-# `update, disable, enable, fix` allows `all` param
+def update_node_id_cache():
+    try:
+        config_manager = ConfigManager()
+        workspace_path = workspace_manager.workspace_path
+
+        cm_cli_path = os.path.join(
+            workspace_path, "custom_nodes", "ComfyUI-Manager", "cm-cli.py"
+        )
+
+        tmp_path = os.path.join(config_manager.get_config_path(), "tmp")
+        if not os.path.exists(tmp_path):
+            os.makedirs(tmp_path)
+
+        cache_path = os.path.join(tmp_path, "node-cache.list")
+        cmd = [sys.executable, cm_cli_path, "export-custom-node-ids", cache_path]
+
+        new_env = os.environ.copy()
+        new_env["COMFYUI_PATH"] = workspace_path
+        subprocess.check_output(cmd, env=new_env)
+    except Exception:
+        pass
 
 
+# `update, disable, enable, fix` allows `all` param
 @app.command(help="Update custom nodes or ComfyUI")
 @tracking.track_command("node")
 def update(
-    args: List[str] = typer.Argument(..., help="update custom nodes"),
+    args: List[str] = typer.Argument(
+        ..., help="update custom nodes", autocompletion=node_completer
+    ),
     channel: Annotated[
         str,
         typer.Option(
@@ -412,11 +458,15 @@ def update(
 
     execute_cm_cli(["update"] + args, channel, mode)
 
+    update_node_id_cache()
+
 
 @app.command(help="Disable custom nodes")
 @tracking.track_command("node")
 def disable(
-    args: List[str] = typer.Argument(..., help="disable custom nodes"),
+    args: List[str] = typer.Argument(
+        ..., help="disable custom nodes", autocompletion=node_completer
+    ),
     channel: Annotated[
         str,
         typer.Option(
@@ -445,7 +495,9 @@ def disable(
 @app.command(help="Enable custom nodes")
 @tracking.track_command("node")
 def enable(
-    args: List[str] = typer.Argument(..., help="enable custom nodes"),
+    args: List[str] = typer.Argument(
+        ..., help="enable custom nodes", autocompletion=node_completer
+    ),
     channel: Annotated[
         str,
         typer.Option(
@@ -475,7 +527,9 @@ def enable(
 @tracking.track_command("node")
 def fix(
     args: List[str] = typer.Argument(
-        ..., help="fix dependencies for specified custom nodes"
+        ...,
+        help="fix dependencies for specified custom nodes",
+        autocompletion=node_completer,
     ),
     channel: Annotated[
         str,

comfy_cli/command/install.py

@@ -8,6 +8,7 @@
 from comfy_cli import constants
 from comfy_cli.constants import GPU_OPTION
 from comfy_cli.workspace_manager import WorkspaceManager, check_comfy_repo
+from comfy_cli.command.custom_nodes.command import update_node_id_cache
 
 
 def get_os_details():
@@ -173,6 +174,8 @@ def execute(
             subprocess.run(["git", "clone", manager_url, manager_repo_dir])
             install_manager_dependencies(repo_dir)
 
+        update_node_id_cache()
+
     os.chdir(repo_dir)
 
     print("")

comfy_cli/utils.py

@@ -57,3 +57,10 @@ def is_running(pid):
         return True
     except psutil.NoSuchProcess:
         return False
+
+
+def create_choice_completer(opts):
+    def f(incomplete: str) -> list[str]:
+        return [opt for opt in opts if opt.startswith(incomplete)]
+
+    return f