-
Notifications
You must be signed in to change notification settings - Fork 4
/
prompt.py
74 lines (55 loc) · 2.08 KB
/
prompt.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
"""Functionality related to prompting the user for some action."""
import logging
import operator
from dataclasses import dataclass
from typing import Callable
import questionary
__all__ = ["PromptChoice", "choice_prompt"]
log = logging.getLogger("moe.cli")
@dataclass
class PromptChoice:
"""A single, user-selectable choice for a CLI prompt.
Attributes:
title: Title of the prompt choice that is displayed to the user.
shortcut_key: Single character the user will use to select the choice.
Important:
Ensure each shortcut key is not in use by another PromptChoice.
func: The function that should get called if a choice is selected.
The definition for how to call ``func`` should be specified by the plugin.
See Also:
`Questionary Choice docs
<https://questionary.readthedocs.io/en/stable/pages/api_reference.html#api-reference>`_
"""
title: str
shortcut_key: str
func: Callable
def choice_prompt(
prompt_choices: list[PromptChoice], question: str = "What do you want to do?"
) -> PromptChoice:
"""Generates a user choice prompt.
Args:
prompt_choices: Prompt choices to be used.
question: Question prompted to the user.
Returns:
The chosen prompt choice.
Raises:
SystemExit: Invalid user input.
"""
prompt_choices.sort(key=operator.attrgetter("shortcut_key"))
questionary_choices: list[questionary.Choice] = []
for prompt_choice in prompt_choices:
questionary_choices.append(
questionary.Choice(
title=prompt_choice.title,
shortcut_key=prompt_choice.shortcut_key,
value=prompt_choice.shortcut_key,
)
)
user_input = questionary.select(
question, choices=questionary_choices, use_shortcuts=True, use_arrow_keys=True
).ask()
for prompt_choice in prompt_choices:
if prompt_choice.shortcut_key == user_input:
return prompt_choice
log.error(f"Invalid option selected. [{user_input=!r}]")
raise SystemExit(1)