From f602660b44646619a89a0323487c1df6d2431c82 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 18 Jul 2024 13:07:32 -0700 Subject: [PATCH 1/2] Fixed: Type ui select inputs with generics. --- comfy_cli/ui.py | 36 +++++++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/comfy_cli/ui.py b/comfy_cli/ui.py index 0c34f51..abbc292 100644 --- a/comfy_cli/ui.py +++ b/comfy_cli/ui.py @@ -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 @@ -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[T]): 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[T]: 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. From 25654f329e0a2d712fa5b608e3209bd95dd5ba6d Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 18 Jul 2024 13:08:32 -0700 Subject: [PATCH 2/2] Fix comment. --- comfy_cli/ui.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_cli/ui.py b/comfy_cli/ui.py index abbc292..c79fc23 100644 --- a/comfy_cli/ui.py +++ b/comfy_cli/ui.py @@ -46,11 +46,11 @@ def prompt_select( Args: question (str): The question to display to the user. - choices (List[T]): A list of 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: - Optional[T]: The selected choice from the user, or None if skipping prompts. + Optional[ChoiceType]: The selected choice from the user, or None if skipping prompts. """ if workspace_manager.skip_prompting and not force_prompting: return None