From 84891699ddd2ce717fb4900348e2ba13516f6973 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Fri, 3 May 2024 00:04:36 +0900
Subject: [PATCH] feat: `comfy update [all|comfy]` command

---
 comfy_cli/cmdline.py | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py
index 4895dc8..0fb5670 100644
--- a/comfy_cli/cmdline.py
+++ b/comfy_cli/cmdline.py
@@ -133,14 +133,31 @@ def install(
     workspace_manager.set_recent_workspace(workspace_path)
 
 
-def update(self):
+@app.command(help="Stop background ComfyUI")
+@tracking.track_command()
+def update(
+    ctx: typer.Context,
+    target: str = typer.Argument("comfy", help="[all|comfy]")
+):
+    if target not in ['all', 'comfy']:
+        typer.echo(
+            f"Invalid target: {target}. Allowed targets are 'all', 'comfy'.",
+            err=True,
+        )
+        raise typer.Exit(code=1)
+
     _env_checker = EnvChecker()
-    print(f"Updating ComfyUI in {self.workspace}...")
-    os.chdir(self.workspace)
-    subprocess.run(["git", "pull"], check=True)
-    subprocess.run(
-        [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], check=True
-    )
+    comfy_path = workspace_manager.get_workspace_comfy_path(ctx)
+
+    if 'all' == target:
+        custom_nodes.command.execute_cm_cli(ctx, ['update', 'all'])
+    else:
+        print(f"Updating ComfyUI in {comfy_path}...")
+        os.chdir(comfy_path)
+        subprocess.run(["git", "pull"], check=True)
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], check=True
+        )
 
 
 # @app.command(help="Run workflow file")
@@ -236,6 +253,7 @@ def launch_comfyui(_env_checker, _config_manager, extra, background=False):
 
 
 @app.command(help="Stop background ComfyUI")
+@tracking.track_command()
 def stop():
     _config_manager = ConfigManager()