mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
ui.prompt: move prompts from exp init to ui.prompt
Some fixes along the way to make it more generic, but it still has lots of features tied with exp init.
This commit is contained in:
Saugat Pachhai (सौगात)
committed by
Saugat Pachhai
Saugat Pachhai
parent
068b8371aa
commit
eac030598c
+32
-117
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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}]: "
|
||||
Reference in New Issue
Block a user