ui: use readline for a more appropriate prompt

beets/ui/__init__.py

@@ -14,7 +14,13 @@ import textwrap
 import traceback
 from typing import TYPE_CHECKING, Any, Literal, TextIO, TypeVar, overload
 
+try:
+    import readline
+except ImportError:
+    readline = None  # type: ignore[assignment]
+
 import confuse
+from rich.ansi import re_ansi
 from rich.logging import RichHandler
 from rich.traceback import install
 from rich_tables.utils import make_console
@@ -65,6 +71,17 @@ class SafeRichHandler(RichHandler):
             self.handleError(record)
 
 
+def _render_prompt(prompt: Any) -> str:
+    with console.capture() as capture:
+        console.print(prompt, end=" ")
+
+    prompt = capture.get()
+    if readline is None:
+        return prompt
+
+    return re_ansi.sub(lambda match: f"\x01{match[0]}\x02", prompt)
+
+
 # Encoding utilities.
 
 
@@ -164,23 +181,14 @@ def should_move(move_opt: bool | None = None) -> bool:
 
 
 def input_(prompt: str | None = None) -> str:
-    """Like `input`, but decodes the result to a Unicode string.
-    Raises a UserError if stdin is not available. The prompt is sent to
-    stdout rather than stderr. A printed between the prompt and the
-    input cursor.
+    """Read a response while keeping styled prompts compatible with readline.
+
+    Raises a user-facing error when input ends before a response is available.
     """
-    # raw_input incorrectly sends prompts to stderr, not stdout, so we
-    # use print_() explicitly to display prompts.
-    # https://bugs.python.org/issue1927
-    if prompt:
-        print_(prompt, end=" ")
-
     try:
-        resp = input()
+        return input(_render_prompt(prompt) if prompt else "")
     except EOFError:
-        raise UserError("stdin stream ended while input required")
-
-    return resp
+        raise UserError("stdin stream ended while input required") from None
 
 
 @overload

test/ui/test_ui_init.py

@@ -4,6 +4,7 @@ import unittest
 from copy import deepcopy
 from pathlib import Path
 from random import random
+from unittest.mock import patch
 
 import pytest
 
@@ -12,6 +13,35 @@ from beets.exceptions import UserError
 from beets.test.helper import BeetsTestCase, IOMixin
 
 
+class TestInput:
+    @pytest.mark.parametrize(
+        "prompt,expected", [(None, ""), ("Prompt:", "Prompt: ")]
+    )
+    def test_passes_prompt_to_input(self, prompt, expected):
+        with patch("builtins.input", return_value="answer") as input_mock:
+            assert ui.input_(prompt) == "answer"
+
+        input_mock.assert_called_once_with(expected)
+
+    @pytest.mark.parametrize(
+        "readline,expected",
+        [
+            (object(), "\x01\x1b[31m\x02Prompt\x01\x1b[0m\x02 "),
+            (None, "\x1b[31mPrompt\x1b[0m "),
+        ],
+    )
+    def test_marks_only_ansi_codes_as_non_printing(self, readline, expected):
+        with patch("beets.ui.readline", readline):
+            assert ui._render_prompt("\x1b[31mPrompt\x1b[0m") == expected
+
+    def test_raises_user_error_at_end_of_input(self):
+        with (
+            patch("builtins.input", side_effect=EOFError),
+            pytest.raises(UserError, match="stdin stream ended"),
+        ):
+            ui.input_("Prompt:")
+
+
 class InputMethodsTest(IOMixin, unittest.TestCase):
     def _print_helper(self, s):
         print(s)