diff --git a/latch_cli/main.py b/latch_cli/main.py index 6b23be3a..06c00999 100644 --- a/latch_cli/main.py +++ b/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()) diff --git a/latch_cli/services/register/register.py b/latch_cli/services/register/register.py index 4f160d11..6ea99fab 100644 --- a/latch_cli/services/register/register.py +++ b/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"), diff --git a/latch_cli/services/workspace.py b/latch_cli/services/workspace.py index e6cb2420..3a821a00 100644 --- a/latch_cli/services/workspace.py +++ b/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