From 958b5ab0096b1f7f8cdb61aaa77e86cff3d2968e Mon Sep 17 00:00:00 2001 From: Ayush Kamat Date: Wed, 24 Apr 2024 14:48:20 -0700 Subject: [PATCH 1/2] add workspace selection to register flow Signed-off-by: Ayush Kamat --- latch_cli/main.py | 10 +++++++ latch_cli/services/register/register.py | 21 ++++++++++++++ latch_cli/services/workspace.py | 37 +++++++++++++------------ 3 files changed, 51 insertions(+), 17 deletions(-) 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..b349ebcd 100644 --- a/latch_cli/services/register/register.py +++ b/latch_cli/services/register/register.py @@ -11,9 +11,11 @@ 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.services.workspace import workspace_tui from ...centromere.ctx import _CentromereCtx from ...centromere.utils import MaybeRemoteDir @@ -273,6 +275,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, @@ -358,6 +361,24 @@ def register( ), "N/A", ) + + if not skip_workspace_selection: + if click.confirm( + f"Workflow will be registered to {ws_name} ({current_workspace()})." + " Use a different one?" + ): + updated = workspace_tui() + if updated is None: + click.secho("Cancelled", bold=True) + return + + user_config.update_workspace(**updated) + else: + click.secho( + "Skipping workspace selection because of --skip-workspace-selection", + bold=True, + ) + 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 From 3f2380c3358de0bb7cfc8ebdd3428e3437d1afef Mon Sep 17 00:00:00 2001 From: Ayush Kamat Date: Wed, 24 Apr 2024 15:20:06 -0700 Subject: [PATCH 2/2] reorder Signed-off-by: Ayush Kamat --- latch_cli/services/register/register.py | 52 ++++++++++++------------- 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/latch_cli/services/register/register.py b/latch_cli/services/register/register.py index b349ebcd..6ea99fab 100644 --- a/latch_cli/services/register/register.py +++ b/latch_cli/services/register/register.py @@ -15,6 +15,7 @@ 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 @@ -333,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, @@ -351,34 +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", - ) - - if not skip_workspace_selection: - if click.confirm( - f"Workflow will be registered to {ws_name} ({current_workspace()})." - " Use a different one?" - ): - updated = workspace_tui() - if updated is None: - click.secho("Cancelled", bold=True) - return - - user_config.update_workspace(**updated) - else: - click.secho( - "Skipping workspace selection because of --skip-workspace-selection", - bold=True, - ) - click.echo( " ".join([ click.style("Target workspace:", fg="bright_blue"),