Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 
 

401 строка
12 KiB

  1. from typing import Any, Generic, List, Optional, TextIO, TypeVar, Union, overload
  2. from . import get_console
  3. from .console import Console
  4. from .text import Text, TextType
  5. PromptType = TypeVar("PromptType")
  6. DefaultType = TypeVar("DefaultType")
  7. class PromptError(Exception):
  8. """Exception base class for prompt related errors."""
  9. class InvalidResponse(PromptError):
  10. """Exception to indicate a response was invalid. Raise this within process_response() to indicate an error
  11. and provide an error message.
  12. Args:
  13. message (Union[str, Text]): Error message.
  14. """
  15. def __init__(self, message: TextType) -> None:
  16. self.message = message
  17. def __rich__(self) -> TextType:
  18. return self.message
  19. class PromptBase(Generic[PromptType]):
  20. """Ask the user for input until a valid response is received. This is the base class, see one of
  21. the concrete classes for examples.
  22. Args:
  23. prompt (TextType, optional): Prompt text. Defaults to "".
  24. console (Console, optional): A Console instance or None to use global console. Defaults to None.
  25. password (bool, optional): Enable password input. Defaults to False.
  26. choices (List[str], optional): A list of valid choices. Defaults to None.
  27. case_sensitive (bool, optional): Matching of choices should be case-sensitive. Defaults to True.
  28. show_default (bool, optional): Show default in prompt. Defaults to True.
  29. show_choices (bool, optional): Show choices in prompt. Defaults to True.
  30. """
  31. response_type: type = str
  32. validate_error_message = "[prompt.invalid]Please enter a valid value"
  33. illegal_choice_message = (
  34. "[prompt.invalid.choice]Please select one of the available options"
  35. )
  36. prompt_suffix = ": "
  37. choices: Optional[List[str]] = None
  38. def __init__(
  39. self,
  40. prompt: TextType = "",
  41. *,
  42. console: Optional[Console] = None,
  43. password: bool = False,
  44. choices: Optional[List[str]] = None,
  45. case_sensitive: bool = True,
  46. show_default: bool = True,
  47. show_choices: bool = True,
  48. ) -> None:
  49. self.console = console or get_console()
  50. self.prompt = (
  51. Text.from_markup(prompt, style="prompt")
  52. if isinstance(prompt, str)
  53. else prompt
  54. )
  55. self.password = password
  56. if choices is not None:
  57. self.choices = choices
  58. self.case_sensitive = case_sensitive
  59. self.show_default = show_default
  60. self.show_choices = show_choices
  61. @classmethod
  62. @overload
  63. def ask(
  64. cls,
  65. prompt: TextType = "",
  66. *,
  67. console: Optional[Console] = None,
  68. password: bool = False,
  69. choices: Optional[List[str]] = None,
  70. case_sensitive: bool = True,
  71. show_default: bool = True,
  72. show_choices: bool = True,
  73. default: DefaultType,
  74. stream: Optional[TextIO] = None,
  75. ) -> Union[DefaultType, PromptType]:
  76. ...
  77. @classmethod
  78. @overload
  79. def ask(
  80. cls,
  81. prompt: TextType = "",
  82. *,
  83. console: Optional[Console] = None,
  84. password: bool = False,
  85. choices: Optional[List[str]] = None,
  86. case_sensitive: bool = True,
  87. show_default: bool = True,
  88. show_choices: bool = True,
  89. stream: Optional[TextIO] = None,
  90. ) -> PromptType:
  91. ...
  92. @classmethod
  93. def ask(
  94. cls,
  95. prompt: TextType = "",
  96. *,
  97. console: Optional[Console] = None,
  98. password: bool = False,
  99. choices: Optional[List[str]] = None,
  100. case_sensitive: bool = True,
  101. show_default: bool = True,
  102. show_choices: bool = True,
  103. default: Any = ...,
  104. stream: Optional[TextIO] = None,
  105. ) -> Any:
  106. """Shortcut to construct and run a prompt loop and return the result.
  107. Example:
  108. >>> filename = Prompt.ask("Enter a filename")
  109. Args:
  110. prompt (TextType, optional): Prompt text. Defaults to "".
  111. console (Console, optional): A Console instance or None to use global console. Defaults to None.
  112. password (bool, optional): Enable password input. Defaults to False.
  113. choices (List[str], optional): A list of valid choices. Defaults to None.
  114. case_sensitive (bool, optional): Matching of choices should be case-sensitive. Defaults to True.
  115. show_default (bool, optional): Show default in prompt. Defaults to True.
  116. show_choices (bool, optional): Show choices in prompt. Defaults to True.
  117. stream (TextIO, optional): Optional text file open for reading to get input. Defaults to None.
  118. """
  119. _prompt = cls(
  120. prompt,
  121. console=console,
  122. password=password,
  123. choices=choices,
  124. case_sensitive=case_sensitive,
  125. show_default=show_default,
  126. show_choices=show_choices,
  127. )
  128. return _prompt(default=default, stream=stream)
  129. def render_default(self, default: DefaultType) -> Text:
  130. """Turn the supplied default in to a Text instance.
  131. Args:
  132. default (DefaultType): Default value.
  133. Returns:
  134. Text: Text containing rendering of default value.
  135. """
  136. return Text(f"({default})", "prompt.default")
  137. def make_prompt(self, default: DefaultType) -> Text:
  138. """Make prompt text.
  139. Args:
  140. default (DefaultType): Default value.
  141. Returns:
  142. Text: Text to display in prompt.
  143. """
  144. prompt = self.prompt.copy()
  145. prompt.end = ""
  146. if self.show_choices and self.choices:
  147. _choices = "/".join(self.choices)
  148. choices = f"[{_choices}]"
  149. prompt.append(" ")
  150. prompt.append(choices, "prompt.choices")
  151. if (
  152. default != ...
  153. and self.show_default
  154. and isinstance(default, (str, self.response_type))
  155. ):
  156. prompt.append(" ")
  157. _default = self.render_default(default)
  158. prompt.append(_default)
  159. prompt.append(self.prompt_suffix)
  160. return prompt
  161. @classmethod
  162. def get_input(
  163. cls,
  164. console: Console,
  165. prompt: TextType,
  166. password: bool,
  167. stream: Optional[TextIO] = None,
  168. ) -> str:
  169. """Get input from user.
  170. Args:
  171. console (Console): Console instance.
  172. prompt (TextType): Prompt text.
  173. password (bool): Enable password entry.
  174. Returns:
  175. str: String from user.
  176. """
  177. return console.input(prompt, password=password, stream=stream)
  178. def check_choice(self, value: str) -> bool:
  179. """Check value is in the list of valid choices.
  180. Args:
  181. value (str): Value entered by user.
  182. Returns:
  183. bool: True if choice was valid, otherwise False.
  184. """
  185. assert self.choices is not None
  186. if self.case_sensitive:
  187. return value.strip() in self.choices
  188. return value.strip().lower() in [choice.lower() for choice in self.choices]
  189. def process_response(self, value: str) -> PromptType:
  190. """Process response from user, convert to prompt type.
  191. Args:
  192. value (str): String typed by user.
  193. Raises:
  194. InvalidResponse: If ``value`` is invalid.
  195. Returns:
  196. PromptType: The value to be returned from ask method.
  197. """
  198. value = value.strip()
  199. try:
  200. return_value: PromptType = self.response_type(value)
  201. except ValueError:
  202. raise InvalidResponse(self.validate_error_message)
  203. if self.choices is not None:
  204. if not self.check_choice(value):
  205. raise InvalidResponse(self.illegal_choice_message)
  206. if not self.case_sensitive:
  207. # return the original choice, not the lower case version
  208. return_value = self.response_type(
  209. self.choices[
  210. [choice.lower() for choice in self.choices].index(value.lower())
  211. ]
  212. )
  213. return return_value
  214. def on_validate_error(self, value: str, error: InvalidResponse) -> None:
  215. """Called to handle validation error.
  216. Args:
  217. value (str): String entered by user.
  218. error (InvalidResponse): Exception instance the initiated the error.
  219. """
  220. self.console.print(error)
  221. def pre_prompt(self) -> None:
  222. """Hook to display something before the prompt."""
  223. @overload
  224. def __call__(self, *, stream: Optional[TextIO] = None) -> PromptType:
  225. ...
  226. @overload
  227. def __call__(
  228. self, *, default: DefaultType, stream: Optional[TextIO] = None
  229. ) -> Union[PromptType, DefaultType]:
  230. ...
  231. def __call__(self, *, default: Any = ..., stream: Optional[TextIO] = None) -> Any:
  232. """Run the prompt loop.
  233. Args:
  234. default (Any, optional): Optional default value.
  235. Returns:
  236. PromptType: Processed value.
  237. """
  238. while True:
  239. self.pre_prompt()
  240. prompt = self.make_prompt(default)
  241. value = self.get_input(self.console, prompt, self.password, stream=stream)
  242. if value == "" and default != ...:
  243. return default
  244. try:
  245. return_value = self.process_response(value)
  246. except InvalidResponse as error:
  247. self.on_validate_error(value, error)
  248. continue
  249. else:
  250. return return_value
  251. class Prompt(PromptBase[str]):
  252. """A prompt that returns a str.
  253. Example:
  254. >>> name = Prompt.ask("Enter your name")
  255. """
  256. response_type = str
  257. class IntPrompt(PromptBase[int]):
  258. """A prompt that returns an integer.
  259. Example:
  260. >>> burrito_count = IntPrompt.ask("How many burritos do you want to order")
  261. """
  262. response_type = int
  263. validate_error_message = "[prompt.invalid]Please enter a valid integer number"
  264. class FloatPrompt(PromptBase[float]):
  265. """A prompt that returns a float.
  266. Example:
  267. >>> temperature = FloatPrompt.ask("Enter desired temperature")
  268. """
  269. response_type = float
  270. validate_error_message = "[prompt.invalid]Please enter a number"
  271. class Confirm(PromptBase[bool]):
  272. """A yes / no confirmation prompt.
  273. Example:
  274. >>> if Confirm.ask("Continue"):
  275. run_job()
  276. """
  277. response_type = bool
  278. validate_error_message = "[prompt.invalid]Please enter Y or N"
  279. choices: List[str] = ["y", "n"]
  280. def render_default(self, default: DefaultType) -> Text:
  281. """Render the default as (y) or (n) rather than True/False."""
  282. yes, no = self.choices
  283. return Text(f"({yes})" if default else f"({no})", style="prompt.default")
  284. def process_response(self, value: str) -> bool:
  285. """Convert choices to a bool."""
  286. value = value.strip().lower()
  287. if value not in self.choices:
  288. raise InvalidResponse(self.validate_error_message)
  289. return value == self.choices[0]
  290. if __name__ == "__main__": # pragma: no cover
  291. from rich import print
  292. if Confirm.ask("Run [i]prompt[/i] tests?", default=True):
  293. while True:
  294. result = IntPrompt.ask(
  295. ":rocket: Enter a number between [b]1[/b] and [b]10[/b]", default=5
  296. )
  297. if result >= 1 and result <= 10:
  298. break
  299. print(":pile_of_poo: [prompt.invalid]Number must be between 1 and 10")
  300. print(f"number={result}")
  301. while True:
  302. password = Prompt.ask(
  303. "Please enter a password [cyan](must be at least 5 characters)",
  304. password=True,
  305. )
  306. if len(password) >= 5:
  307. break
  308. print("[prompt.invalid]password too short")
  309. print(f"password={password!r}")
  310. fruit = Prompt.ask("Enter a fruit", choices=["apple", "orange", "pear"])
  311. print(f"fruit={fruit!r}")
  312. doggie = Prompt.ask(
  313. "What's the best Dog? (Case INSENSITIVE)",
  314. choices=["Border Terrier", "Collie", "Labradoodle"],
  315. case_sensitive=False,
  316. )
  317. print(f"doggie={doggie!r}")
  318. else:
  319. print("[b]OK :loudly_crying_face:")