Fix #15: ask_* functions prompt can be colorized

Author: dmerejkowsky

README.rst

@@ -61,6 +61,6 @@ Example:
     with_sugar = cli_ui.ask_yes_no("With sugar?", default=False)
 
     fruits = ["apple", "orange", "banana"]
-    selected_fruit = cli_ui.ask_choice("Choose a fruit", fruits)
+    selected_fruit = cli_ui.ask_choice("Choose a fruit", choices=fruits)
 
     #  ... and more!

cli_ui/__init__.py

@@ -449,22 +449,28 @@ def read_password() -> str:
     return getpass.getpass(prompt="")
 
 
-def ask_string(question: str, default: Optional[str] = None) -> Optional[str]:
+def get_ask_tokens(tokens: Sequence[Token]) -> List[Token]:
+    return [green, "::", reset] + list(tokens) + [reset]  # type: ignore
+
+
+def ask_string(*question: Token, default: Optional[str] = None) -> Optional[str]:
     """Ask the user to enter a string.
     """
+    tokens = get_ask_tokens(question)
     if default:
-        question += " (%s)" % default
-    info(green, "::", reset, question)
+        tokens.append("(%s)" % default)
+    info(*tokens)
     answer = read_input()
     if not answer:
         return default
     return answer
 
 
-def ask_password(question: str) -> str:
+def ask_password(*question: Token) -> str:
     """Ask the user to enter a password.
     """
-    info(green, "::", reset, question)
+    tokens = get_ask_tokens(question)
+    info(*tokens)
     answer = read_password()
     return answer
 
@@ -473,7 +479,7 @@ def ask_password(question: str) -> str:
 
 
 def ask_choice(
-    input_text: str, choices: List[Any], *, func_desc: Optional[FuncDesc] = None
+    *prompt: Token, choices: List[Any], func_desc: Optional[FuncDesc] = None
 ) -> Any:
     """Ask the user to choose from a list of choices.
 
@@ -492,7 +498,8 @@ def ask_choice(
     """
     if func_desc is None:
         func_desc = lambda x: str(x)
-    info(green, "::", reset, input_text)
+    tokens = get_ask_tokens(prompt)
+    info(*tokens)
     choices.sort(key=func_desc)
     for i, choice in enumerate(choices, start=1):
         choice_desc = func_desc(choice)
@@ -517,13 +524,15 @@ def ask_choice(
     return res
 
 
-def ask_yes_no(question: str, default: bool = False) -> bool:
+def ask_yes_no(*question: Token, default: bool = False) -> bool:
     """Ask the user to answer by yes or no"""
     while True:
+        tokens = [green, "::", reset] + list(question) + [reset]
         if default:
-            info(green, "::", reset, question, "(Y/n)")
+            tokens.append("(Y/n)")
         else:
-            info(green, "::", reset, question, "(y/N)")
+            tokens.append("(y/N)")
+        info(*tokens)  # type: ignore
         answer = read_input()
         if answer.lower() in ["y", "yes"]:
             return True
@@ -639,7 +648,7 @@ def new_message(*args: Any, **kwargs: Any) -> None:
     # stop monkey patching
     message = old_message
     fruits = ["apple", "orange", "banana"]
-    answer = ask_choice("Choose a fruit", fruits)
+    answer = ask_choice("Choose a fruit", choices=fruits)
     info("You chose:", answer)
 
 

cli_ui/tests/test_cli_ui.py

@@ -159,6 +159,15 @@ def test_ask_string() -> None:
         assert res == "milk"
 
 
+def test_ask_colored_message() -> None:
+    with mock.patch("builtins.input") as m:
+        m.side_effect = ["y"]
+        res = cli_ui.ask_yes_no(
+            "Deploy to", cli_ui.bold, "prod", cli_ui.reset, "?", default=False
+        )
+        assert res
+
+
 def test_ask_password() -> None:
     with mock.patch("getpass.getpass") as m:
         m.side_effect = ["chocolate!", ""]
@@ -214,7 +223,7 @@ def func_desc(fruit: Fruit) -> str:
     with mock.patch("builtins.input") as m:
         m.side_effect = ["nan", "5", "2"]
         actual = cli_ui.ask_choice(
-            "Select a fruit", fruits, func_desc=operator.attrgetter("name")
+            "Select a fruit", choices=fruits, func_desc=operator.attrgetter("name")
         )
         assert actual.name == "banana"
         assert actual.price == 10
@@ -224,15 +233,15 @@ def func_desc(fruit: Fruit) -> str:
 def test_ask_choice_empty_input() -> None:
     with mock.patch("builtins.input") as m:
         m.side_effect = [""]
-        res = cli_ui.ask_choice("Select a animal", ["cat", "dog", "cow"])
+        res = cli_ui.ask_choice("Select a animal", choices=["cat", "dog", "cow"])
         assert res is None
 
 
 def test_ask_choice_ctrl_c() -> None:
     with pytest.raises(KeyboardInterrupt):
         with mock.patch("builtins.input") as m:
             m.side_effect = KeyboardInterrupt
-            cli_ui.ask_choice("Select a animal", ["cat", "dog", "cow"])
+            cli_ui.ask_choice("Select a animal", choices=["cat", "dog", "cow"])
 
 
 def test_quiet(message_recorder: MessageRecorder) -> None: