add workspace selection to register flow

latch_cli/main.py

@@ -387,6 +387,14 @@ def execute(
     type=bool,
     help="Skip the confirmation dialog.",
 )
+@click.option(
+    "-w",
+    "--skip-workspace-selection",
+    is_flag=True,
+    default=False,
+    type=bool,
+    help="Skip workspace selection.",
+)
 @click.option(
     "--open",
     "-o",
@@ -419,6 +427,7 @@ def register(
     remote: bool,
     docker_progress: str,
     yes: bool,
+    skip_workspace_selection: bool,
     open: bool,
     snakefile: Optional[Path],
     cache_tasks: bool,
@@ -440,6 +449,7 @@ def register(
         disable_auto_version=disable_auto_version,
         remote=remote,
         skip_confirmation=yes,
+        skip_workspace_selection=skip_workspace_selection,
         open=open,
         snakefile=snakefile,
         progress_plain=(docker_progress == "auto" and not sys.stdout.isatty())

latch_cli/services/register/register.py

@@ -11,9 +11,12 @@
 import click
 import gql
 import latch_sdk_gql.execute as l_gql
+from latch_sdk_config.user import user_config
 from scp import SCPClient
 
 from latch.utils import current_workspace, get_workspaces
+from latch_cli.click_utils import bold, color
+from latch_cli.services.workspace import workspace_tui
 
 from ...centromere.ctx import _CentromereCtx
 from ...centromere.utils import MaybeRemoteDir
@@ -273,6 +276,7 @@ def register(
     remote: bool = False,
     open: bool = False,
     skip_confirmation: bool = False,
+    skip_workspace_selection: bool = False,
     snakefile: Optional[Path] = None,
     progress_plain: bool = False,
     cache_tasks: bool = False,
@@ -330,6 +334,29 @@ def register(
             click.secho("\n`snakemake` package is not installed.", fg="red", bold=True)
             sys.exit(1)
 
+    workspaces = get_workspaces()
+    ws_name = next(
+        (
+            x[1]
+            for x in workspaces.items()
+            if x[0] == current_workspace()
+            or (current_workspace() == "" and x[1] == "Personal Workspace")
+        ),
+        "N/A",
+    )
+
+    if not skip_workspace_selection:
+        if click.confirm(
+            f"Workflow will be registered to {ws_name} ({current_workspace()})."
+            " Use a different workspace?"
+        ):
+            updated = workspace_tui()
+            if updated is None:
+                click.secho("Cancelled", bold=True)
+                return
+
+            user_config.update_workspace(**updated)
+
     with _CentromereCtx(
         Path(pkg_root),
         disable_auto_version=disable_auto_version,
@@ -348,16 +375,6 @@ def register(
         )
         click.echo(" ".join([click.style("Version:", fg="bright_blue"), ctx.version]))
 
-        workspaces = get_workspaces()
-        ws_name = next(
-            (
-                x[1]
-                for x in workspaces.items()
-                if x[0] == current_workspace()
-                or (current_workspace() == "" and x[1] == "Personal Workspace")
-            ),
-            "N/A",
-        )
         click.echo(
             " ".join([
                 click.style("Target workspace:", fg="bright_blue"),

latch_cli/services/workspace.py

@@ -1,4 +1,4 @@
-from typing import List, TypedDict
+from typing import List, Optional, TypedDict
 
 import click
 from latch_sdk_config.user import user_config
@@ -12,12 +12,7 @@ class WSInfo(TypedDict):
     name: str
 
 
-def workspace():
-    """Opens a terminal user interface in which a user can select the workspace
-    the want to switch to.
-
-    Like `get_executions`, this function should only be called from the CLI.
-    """
+def workspace_tui() -> Optional[WSInfo]:
     data = get_workspaces()
 
     old_id = current_workspace()
@@ -32,22 +27,30 @@ def workspace():
         if id == old_id:
             display_name = f"{name}{selected_marker}"
 
-        options.append(
-            {
-                "display_name": display_name,
-                "value": {
-                    "workspace_id": id,
-                    "name": name,
-                },
-            }
-        )
+        options.append({
+            "display_name": display_name,
+            "value": {
+                "workspace_id": id,
+                "name": name,
+            },
+        })
 
-    selected_option = select_tui(
+    return select_tui(
         title="Select Workspace",
         options=options,
         clear_terminal=False,
     )
 
+
+def workspace():
+    """Opens a terminal user interface in which a user can select the workspace
+    the want to switch to.
+
+    Like `get_executions`, this function should only be called from the CLI.
+    """
+    old_id = current_workspace()
+
+    selected_option = workspace_tui()
     if selected_option is None:
         return