import os
import sys
import termios
import tty
import webbrowser
from pathlib import Path
from typing import List
from flytekit.core.workflow import PythonFunctionWorkflow
from google.protobuf.json_format import MessageToJson
from latch_sdk_config.latch import config
import latch_cli.menus as menus
from latch.utils import current_workspace, retrieve_or_login
from latch_cli.centromere.utils import _import_flyte_objects
from latch_cli.tinyrequests import post
# TODO(ayush): make this import the `wf` directory and use the package root
# instead of the workflow name. also redo the frontend, also make it open the
# page
[docs]def preview(pkg_root: Path):
"""Generate a preview of the parameter interface for a workflow.
This will allow a user to see how their parameter interface will look
without having to first register their workflow.
Args:
pkg_root: A valid path pointing to the worklow code a user wishes to
preview. The path can be absolute or relative.
Example:
>>> preview("wf.__init__.alphafold_wf")
"""
try:
modules = _import_flyte_objects([pkg_root.resolve()])
wfs: dict[str, PythonFunctionWorkflow] = {}
for module in modules:
for flyte_obj in module.__dict__.values():
if isinstance(flyte_obj, PythonFunctionWorkflow):
wfs[flyte_obj.name] = flyte_obj
if len(wfs) == 0:
raise ValueError(f"Unable to find a workflow definition in {pkg_root}")
except ImportError as e:
raise ValueError(
f"Unable to find {e.name} - make sure that all necessary packages"
" are installed and you have the correct function name."
)
wf = list(wfs.values())[0]
if len(wfs) > 1:
wf = wfs[
_select_workflow_tui(
title="Select which workflow to preview",
options=list(wfs.keys()),
)
]
resp = post(
url=config.api.workflow.preview,
headers={"Authorization": f"Bearer {retrieve_or_login()}"},
json={
"workflow_ui_preview": MessageToJson(wf.interface.to_flyte_idl().inputs),
"ws_account_id": current_workspace(),
},
)
resp.raise_for_status()
url = f"{config.console_url}/preview/parameters"
webbrowser.open(url)
# TODO(ayush): abstract this logic in a unified interface that all tui commands use
def _select_workflow_tui(title: str, options: List[str], clear_terminal: bool = True):
"""
Renders a terminal UI that allows users to select one of the options
listed in `options`
Args:
title: The title of the selection window.
options: A list of names for each of the options.
clear_terminal: Whether or not to clear the entire terminal window
before displaying - default False
"""
if len(options) == 0:
raise ValueError("No options given")
def render(
curr_selected: int,
start_index: int = 0,
max_per_page: int = 10,
indent: str = " ",
) -> int:
if curr_selected < 0 or curr_selected >= len(options):
curr_selected = 0
menus._print(title)
menus.line_down(2)
num_lines_rendered = 4 # 4 "extra" lines for header + footer
for i in range(start_index, start_index + max_per_page):
if i >= len(options):
break
name = options[i]
if i == curr_selected:
color = "\x1b[38;5;40m"
bold = "\x1b[1m"
reset = "\x1b[0m"
menus._print(f"{indent}{color}{bold}{name}{reset}\x1b[1E")
else:
menus._print(f"{indent}{name}\x1b[1E")
num_lines_rendered += 1
menus.line_down(1)
control_str = "[ARROW-KEYS] Navigate\t[ENTER] Select\t[Q] Quit"
menus._print(control_str)
menus.line_up(num_lines_rendered - 1)
menus._show()
return num_lines_rendered
old_settings = termios.tcgetattr(sys.stdin.fileno())
tty.setraw(sys.stdin.fileno())
curr_selected = 0
start_index = 0
_, term_height = os.get_terminal_size()
menus.remove_cursor()
if not clear_terminal:
_, curs_height = menus.current_cursor_position()
max_per_page = term_height - curs_height - 4
else:
menus.clear_screen()
menus.move_cursor((0, 0))
max_per_page = term_height - 4
num_lines_rendered = render(
curr_selected,
start_index=start_index,
max_per_page=max_per_page,
)
try:
while True:
b = menus.read_bytes(1)
if b == b"\r":
return options[curr_selected]
elif b == b"\x1b":
b = menus.read_bytes(2)
if b == b"[A": # Up Arrow
curr_selected = max(curr_selected - 1, 0)
if (
curr_selected - start_index < max_per_page // 2
and start_index > 0
):
start_index -= 1
elif b == b"[B": # Down Arrow
curr_selected = min(curr_selected + 1, len(options) - 1)
if (
curr_selected - start_index > max_per_page // 2
and start_index < len(options) - max_per_page
):
start_index += 1
else:
continue
menus.clear(num_lines_rendered)
num_lines_rendered = render(
curr_selected,
start_index=start_index,
max_per_page=max_per_page,
)
except KeyboardInterrupt: ...
finally:
menus.clear(num_lines_rendered)
menus.reveal_cursor()
menus._show()
termios.tcsetattr(sys.stdin.fileno(), termios.TCSANOW, old_settings)