diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index 4df83e20e..15b4e4115 100644 --- a/dvc/repo/experiments/init.py +++ b/dvc/repo/experiments/init.py @@ -1,6 +1,7 @@ import logging import os from contextlib import contextmanager +from functools import partial from typing import ( TYPE_CHECKING, Any, @@ -8,19 +9,18 @@ from typing import ( Dict, Iterable, List, - Optional, + Tuple, + Union, cast, ) from funcy import compact, lremove -from rich.prompt import Confirm, Prompt from rich.rule import Rule from rich.syntax import Syntax from dvc.exceptions import DvcException from dvc.stage import PipelineStage from dvc.stage.serialize import to_pipeline_file -from dvc.types import OptStr from dvc.utils.serialize import dumps_yaml if TYPE_CHECKING: @@ -42,113 +42,25 @@ PROMPTS = { } -class RetryPrompt(Exception): - """Used for signalling whether prompts should be retried again.""" - - -class RichInputMixin: - """Prevents exc message from printing in the same line on Ctrl + D/C.""" - - @classmethod - def get_input(cls, *args, **kwargs) -> str: - try: - return super().get_input(*args, **kwargs) # type: ignore[misc] - except (KeyboardInterrupt, EOFError): - ui.error_write() - raise - - -class ConfirmationPrompt(RichInputMixin, Confirm): - def make_prompt(self, default): - prompt = self.prompt.copy() - prompt.end = "" - - prompt.append(" [") - yes, no = self.choices - for idx, (val, choice) in enumerate(((True, yes), (False, no))): - if idx: - prompt.append("/") - if val == default: - prompt.append(choice.upper(), "green") - else: - prompt.append(choice) - prompt.append("] ") - return prompt - - -class RequiredPrompt(RichInputMixin, Prompt): - def process_response(self, value: str): - from rich.prompt import InvalidResponse - - ret = super().process_response(value) - if not ret: - raise InvalidResponse( - "[prompt.invalid]Response required. Please try again." - ) - return ret - - def render_default(self, default): - from rich.text import Text - - return Text(f"{default!s}", "green") - - -class SkippablePrompt(RequiredPrompt): - skip_value: str = "n" - - def process_response(self, value: str): - ret = super().process_response(value) - return None if ret == self.skip_value else ret - - def make_prompt(self, default): - prompt = self.prompt.copy() - prompt.end = "" - - prompt.append(" [") - if ( - default is not ... - and self.show_default - and isinstance(default, (str, self.response_type)) - ): - _default = self.render_default(default) - prompt.append(_default) - prompt.append(", ") - - prompt.append(f"{self.skip_value} to omit", style="italic") - prompt.append("]") - prompt.append(self.prompt_suffix) - return prompt - - -def _prompt( - key: str, - validator: Optional[Callable[[str, Any], None]] = None, - default: OptStr = None, -) -> str: - prompt_cls = RequiredPrompt if key == "cmd" else SkippablePrompt - kwargs = {"default": default} if default is not None else {} - while True: - value = prompt_cls.ask( # type: ignore[call-overload] - PROMPTS[key], console=ui.error_console, **kwargs - ) - if validator: - try: - validator(key, value) - except RetryPrompt as exc: - ui.error_write(f"[red]{exc}[/]", styled=True) - continue - return value - - def _prompts( keys: Iterable[str], defaults: Dict[str, str], - validator: Callable[[str, Any], None] = None, + validator: Callable[[str, str], Union[str, Tuple[str, str]]] = None, ) -> Dict[str, str]: - return { - key: _prompt(key, default=defaults.get(key), validator=validator) - for key in keys - } + from dvc.ui.prompt import Prompt + + ret: Dict[str, str] = {} + for key in keys: + value = Prompt.prompt_( + PROMPTS[key], + console=ui.error_console, + default=defaults.get(key), + validator=partial(validator, key) if validator else None, + allow_omission=key != "cmd", + ) + if value: + ret[key] = value + return ret @contextmanager @@ -172,7 +84,7 @@ def init_interactive( name: str, defaults: Dict[str, str], provided: Dict[str, str], - validator: Callable[[str, Any], None] = None, + validator: Callable[[str, str], Union[str, Tuple[str, str]]] = None, show_tree: bool = False, live: bool = False, ) -> Dict[str, str]: @@ -225,7 +137,7 @@ def init_interactive( ui.error_write() ret.update(_prompts(primary, defaults, validator=validator)) ret.update(_prompts(secondary, defaults, validator=validator)) - return compact(ret) + return ret def _check_stage_exists( @@ -268,9 +180,10 @@ def init( with_live = type == "live" - def validate_prompts_input(key: str, value: Any) -> None: - if value is None: - return + def validate_prompts_input( + key: str, value: str + ) -> Union[Any, Tuple[Any, str]]: + from dvc.ui.prompt import InvalidResponse if key == "params": assert isinstance(value, str) @@ -280,17 +193,17 @@ def init( reason = "does not exist" if isinstance(exc, IsADirectoryError): reason = "is a directory" - raise RetryPrompt( - f"'{value}' {reason}. " + raise InvalidResponse( + f"[prompt.invalid]'{value}' {reason}. " "Please retry with an existing parameters file." ) elif key in ("code", "data"): if not os.path.exists(value): - ui.error_write( + return value, ( f"[yellow]'{value}' does not exist in the workspace. " - '"exp run" may fail.[/]', - styled=True, + '"exp run" may fail.[/]' ) + return value if interactive: defaults = init_interactive( @@ -339,7 +252,9 @@ def init( syn = Syntax(_yaml, "yaml", theme="ansi_dark") ui.error_write(syn, styled=True) - if not interactive or ConfirmationPrompt.ask( + from dvc.ui.prompt import Confirm + + if not interactive or Confirm.ask( "Do you want to add the above contents to dvc.yaml?", console=ui.error_console, default=True, diff --git a/dvc/ui/prompt.py b/dvc/ui/prompt.py new file mode 100644 index 000000000..7943584b8 --- /dev/null +++ b/dvc/ui/prompt.py @@ -0,0 +1,185 @@ +from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union + +import rich.prompt + +if TYPE_CHECKING: + from rich.console import Console as RichConsole + + from . import RichText + + +TextType = Union["RichText", str] + + +class InvalidResponse(rich.prompt.InvalidResponse): + pass + + +class RichInputMixin: + """Prevents exc message from printing in the same line on Ctrl + D/C.""" + + @classmethod + def get_input(cls, console: "RichConsole", *args, **kwargs) -> str: + try: + return super().get_input( # type: ignore[misc] + console, *args, **kwargs + ) + except (KeyboardInterrupt, EOFError): + console.print() + raise + + +class Prompt(RichInputMixin, rich.prompt.Prompt): + """Extended version supports custom validations and omission of values.""" + + omit_value: str = "n" + + def __init__( + self, + *args: Any, + allow_omission: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + self.allow_omission: bool = allow_omission + + @classmethod + def ask( # pylint: disable=arguments-differ + cls, + *args: Any, + allow_omission: bool = False, + validator: Optional[ + Callable[[str], Union[str, Tuple[str, str]]] + ] = None, + **kwargs: Any, + ) -> str: + """Extended to pass validator argument. + + Args: + allow_omission: Allows omitting prompts + validator: Validates entered value through prompts + validator can raise an InvalidResponseError, after which + it'll reprompt user. It can also return: + - processed/modified value that will be return to the user. + - tuple of processed value and a message to show to the user. + Note that the validator is always required to return a value, + even if it does not process/modify it. + """ + default = kwargs.pop("default", None) + stream = kwargs.pop("stream", None) + + return cls(*args, allow_omission=allow_omission, **kwargs)( + default=default if default is not None else ..., + stream=stream, + validator=validator, + ) + + @classmethod + def prompt_( + cls, + *args: Any, + allow_omission: bool = False, + validator: Optional[ + Callable[[str], Union[str, Tuple[str, str]]] + ] = None, + **kwargs: Any, + ) -> Optional[str]: + """Extends `.ask()` to allow skipping/returning None type.""" + value = cls.ask( + *args, allow_omission=allow_omission, validator=validator, **kwargs + ) + if allow_omission and value == cls.omit_value: + return None + return value + + def __call__( + self, + *args: Any, + validator: Optional[ + Callable[[str], Union[str, Tuple[str, str]]] + ] = None, + **kwargs: Any, + ) -> str: + """Supports validating response and show warning message.""" + + # We cannot put this in `process_response` as it does not validate + # `default` values. + while True: + value = super().__call__(*args, **kwargs) + if validator is None or ( + self.allow_omission and value == self.omit_value + ): + return value + + try: + validated = validator(value) + except InvalidResponse as exc: + self.on_validate_error(value, exc) + continue + else: + if isinstance(validated, tuple): + value, message = validated + else: + value, message = validated, "" + if message: + self.console.print(message) + return value + + def process_response(self, value: str) -> str: + """Disallow empty values.""" + ret = super().process_response(value) + if not ret: + raise InvalidResponse( + "[prompt.invalid]Response required. Please try again." + ) + return ret + + def render_default(self, default): + from rich.text import Text + + return Text(f"{default!s}", "green") + + def make_prompt(self, default): + prompt = self.prompt.copy() + prompt.end = "" + + parts = [] + if ( + default is not ... + and self.show_default + and isinstance(default, (str, self.response_type)) + ): + _default = self.render_default(default) + parts.append(_default) + + if self.allow_omission and parts: + from rich.text import Text + + parts.append(Text(f", {self.omit_value} to omit", style="italic")) + + if parts: + parts = [" [", *parts, "]"] + + for part in parts: + prompt.append(part) + prompt.append(self.prompt_suffix) + return prompt + + +class Confirm(RichInputMixin, rich.prompt.Confirm): + def make_prompt(self, default): + prompt = self.prompt.copy() + prompt.end = "" + + prompt.append(" [") + yes, no = self.choices + for idx, (val, choice) in enumerate(((True, yes), (False, no))): + if idx: + prompt.append("/") + if val == default: + prompt.append(choice.upper(), "green") + else: + prompt.append(choice) + prompt.append("]") + prompt.append(self.prompt_suffix) + return prompt diff --git a/tests/unit/ui/test_prompt.py b/tests/unit/ui/test_prompt.py new file mode 100644 index 000000000..9a0f6c9ad --- /dev/null +++ b/tests/unit/ui/test_prompt.py @@ -0,0 +1,146 @@ +import io +from unittest.mock import call + +import pytest +from rich.console import Console as RichConsole + +from dvc.ui.prompt import Confirm, InvalidResponse, Prompt + +# pylint: disable=no-member + + +@pytest.mark.parametrize("side_effect", [EOFError, KeyboardInterrupt]) +@pytest.mark.parametrize("prompt_cls", [Prompt, Confirm]) +def test_prompt_interrupt_message_not_interleaved( + prompt_cls, side_effect, mocker +): + """Test that the interrupt message is not interleaved with the prompt.""" + prompt = "should print exception in the next line" + console = RichConsole(file=io.StringIO()) + + with pytest.raises(side_effect): + mocker.patch.object(console, "input", side_effect=side_effect) + prompt_cls.ask(prompt=prompt, console=console) + # as we are raising side_effect on the input, unfortunately we cannot + # assert for the prompt string, only the newline. + assert console.file.getvalue() == "\n" + + +def test_prompt_str(): + console = RichConsole(file=io.StringIO()) + assert ( + Prompt.ask( + "What is your name?", console=console, stream=io.StringIO("foo") + ) + == "foo" + ) + assert console.file.getvalue() == "What is your name?: " + + +def test_prompt_disallows_empty_response(mocker): + console = RichConsole(file=io.StringIO()) + Prompt.prompt_( + "What is your name?", + console=console, + stream=io.StringIO("\nfoo"), + ) + expected = ( + "What is your name?: Response required. Please try again.\n" + "What is your name?: " + ) + assert console.file.getvalue() == expected + + +def test_prompt_with_conditional_skip(mocker): + console = RichConsole(file=io.StringIO()) + mocked_validator = mocker.MagicMock(return_value="something") + assert ( + Prompt.prompt_( + "What is your name?", + console=console, + allow_omission=True, + stream=io.StringIO("n\n"), + default="default", + validator=mocked_validator, + ) + is None + ) + assert ( + console.file.getvalue() == "What is your name? [default, n to omit]: " + ) + mocked_validator.assert_not_called() + + +def test_prompt_retries_on_invalid_response_from_validator(mocker): + console = RichConsole(file=io.StringIO()) + mocked_validator = mocker.MagicMock( + side_effect=[InvalidResponse("it is a number"), "foo"] + ) + assert ( + Prompt.ask( + "what is your name?", + console=console, + stream=io.StringIO("3\nfoo"), + validator=mocked_validator, + ) + == "foo" + ) + mocked_validator.assert_has_calls([call("3"), call("foo")]) + + expected = "what is your name?: it is a number\nwhat is your name?: " + assert console.file.getvalue() == expected + + +def test_prompt_shows_message_from_validator_response(mocker): + console = RichConsole(file=io.StringIO()) + mocked_validator = mocker.MagicMock( + return_value=("foo@email", "failed to send a verification email") + ) + assert ( + Prompt.ask( + "what is your email?", + console=console, + stream=io.StringIO("foo@email"), + validator=mocked_validator, + ) + == "foo@email" + ) + mocked_validator.assert_has_calls([call("foo@email")]) + + expected = "what is your email?: failed to send a verification email\n" + assert console.file.getvalue() == expected + + +def test_prompt_default(mocker): + console = RichConsole(file=io.StringIO()) + mocked_validator = mocker.MagicMock(return_value="default") + assert ( + Prompt.ask( + "What is your name?", + console=console, + stream=io.StringIO(""), + default="default", + validator=mocked_validator, + ) + == "default" + ) + assert console.file.getvalue() == "What is your name? [default]: " + mocked_validator.assert_has_calls([call("default")]) + + +@pytest.mark.parametrize( + "default, inner", [(True, "Y/n"), (False, "y/N"), (Ellipsis, "y/n")] +) +def test_prompt_confirm(default, inner): + console = RichConsole(file=io.StringIO()) + assert ( + Confirm.ask( + "exit?", + console=console, + stream=io.StringIO("y\n"), + default=default, + ) + is True + ) + output = console.file.getvalue() + assert output == f"exit? [{inner}]: "