Skip to content

Commit

Permalink
Fixed: Type ui select inputs with generics. (#134)
Browse files Browse the repository at this point in the history
* Fixed: Type ui select inputs with generics.

* Fix comment.
  • Loading branch information
robinjhuang authored Jul 18, 2024
1 parent c79a1f6 commit 8d6f096
Showing 1 changed file with 25 additions and 11 deletions.
36 changes: 25 additions & 11 deletions comfy_cli/ui.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from typing import List, Tuple

from typing import List, Tuple, Optional, TypeVar, Union, Dict, Any
from enum import Enum
from questionary import Choice
import questionary
import typer
from rich.console import Console
Expand Down Expand Up @@ -33,47 +34,60 @@ def show_progress(iterable, total, description="Downloading..."):
progress.update(task, advance=len(chunk))


def prompt_select(question: str, choices: list, force_prompting: bool = False) -> str:
ChoiceType = Union[str, Choice, Dict[str, Any]]

def prompt_select(
question: str,
choices: List[ChoiceType],
force_prompting: bool = False
) -> Optional[ChoiceType]:
"""
Asks a single select question using questionary and returns the selected response.
Args:
question (str): The question to display to the user.
choices (list): A list of string choices for the user to select from.
choices (List[ChoiceType]): A list of choices for the user to select from.
force_prompting (bool): Whether to force prompting even if skip_prompting is set.
Returns:
str: The selected choice from the user.
Optional[ChoiceType]: The selected choice from the user, or None if skipping prompts.
"""
if workspace_manager.skip_prompting and not force_prompting:
return None
return questionary.select(question, choices=choices).ask()


E = TypeVar('E', bound=Enum)

def prompt_select_enum(
question: str, choices: list, force_prompting: bool = False
) -> str:
question: str,
choices: List[E],
force_prompting: bool = False
) -> Optional[E]:
"""
Asks a single select question using questionary and returns the selected response.
Args:
question (str): The question to display to the user.
choices (list): A list of Enum choices for the user to select from.
choices (List[E]): A list of Enum choices for the user to select from.
force_prompting (bool): Whether to force prompting even if skip_prompting is set.
Returns:
str: The selected choice from the user.
Optional[E]: The selected Enum choice from the user, or None if skipping prompts.
"""
if workspace_manager.skip_prompting and not force_prompting:
return None

choice_map = {choice.value: choice for choice in choices}
display_choices = list(choice_map.keys())

selected = questionary.select(question, choices=display_choices).ask()

return choice_map[selected]
return choice_map[selected] if selected is not None else None


def prompt_input(
question: str, default: str = None, force_prompting: bool = False
question: str, default: str = "", force_prompting: bool = False
) -> str:
"""
Asks the user for an input using questionary.
Expand Down

0 comments on commit 8d6f096

Please sign in to comment.