From 63dba45c142c4fc7a6d2a65b323ae987246ec840 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Fri, 28 May 2021 09:30:56 +0545 Subject: [PATCH] ui: write gets a styled param for rich console printing (#6037) --- dvc/command/checkout.py | 19 +++++++++---------- dvc/ui/__init__.py | 25 +++++++++++++++++-------- tests/func/test_checkout.py | 6 ++++-- tests/func/test_data_cloud.py | 6 ++---- tests/unit/command/test_checkout.py | 5 ++--- 5 files changed, 34 insertions(+), 27 deletions(-) diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py index 7981cdc41..4938c6761 100644 --- a/dvc/command/checkout.py +++ b/dvc/command/checkout.py @@ -1,8 +1,6 @@ import argparse import operator -import colorama - from dvc.command import completion from dvc.command.base import CmdBase, append_doc_link from dvc.exceptions import CheckoutError @@ -10,21 +8,22 @@ from dvc.ui import ui def log_changes(stats): - colors = [ - ("modified", colorama.Fore.YELLOW), - ("added", colorama.Fore.GREEN), - ("deleted", colorama.Fore.RED), - ] + colors = { + "modified": "yellow", + "added": "green", + "deleted": "red", + } - for state, color in colors: + for state, color in colors.items(): entries = stats.get(state) if not entries: continue - nc = colorama.Fore.RESET for entry in entries: - ui.write(f"{color}{state[0].upper()}{nc}", entry, sep="\t") + ui.write( + f"[{color}]{state[0].upper()}", entry, styled=True, sep="\t" + ) class CmdCheckout(CmdBase): diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py index 3bf428470..3a06ae0ef 100644 --- a/dvc/ui/__init__.py +++ b/dvc/ui/__init__.py @@ -7,7 +7,6 @@ from typing import ( Iterable, Optional, Sequence, - TextIO, Union, ) @@ -58,15 +57,13 @@ class Console: style: str = None, sep: str = None, end: str = None, - flush: bool = False, ) -> None: return self.write( *objects, style=style, sep=sep, end=end, - file=sys.stderr, - flush=flush, + stderr=True, ) def write( @@ -75,16 +72,22 @@ class Console: style: str = None, sep: str = None, end: str = None, - file: TextIO = None, - flush: bool = False, + stderr: bool = False, force: bool = False, + styled: bool = False, ) -> None: + sep = " " if sep is None else sep + end = "\n" if end is None else end if not self._enabled and not force: return - file = file or sys.stdout + if styled: + console = self.error_console if stderr else self.rich_console + return console.print(*objects, sep=sep, end=end) + + file = sys.stderr if stderr else sys.stdout values = (self.formatter.format(obj, style=style) for obj in objects) - return print(*values, sep=sep, end=end, file=file, flush=flush) + return print(*values, sep=sep, end=end, file=file) @staticmethod def progress(*args, **kwargs) -> Tqdm: @@ -132,6 +135,12 @@ class Console: return console.Console() + @cached_property + def error_console(self): + from rich import console + + return console.Console(stderr=True) + def table( self, data: "TableData", diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py index d2e373fad..3501d43ea 100644 --- a/tests/func/test_checkout.py +++ b/tests/func/test_checkout.py @@ -26,7 +26,6 @@ from dvc.utils import relpath from dvc.utils.fs import remove, walk_files from dvc.utils.serialize import dump_yaml, load_yaml from tests.basic_env import TestDvc, TestDvcGit -from tests.func.parsing.test_errors import escape_ansi from tests.func.test_repro import TestRepro logger = logging.getLogger("dvc") @@ -688,7 +687,10 @@ def test_stats_on_show_changes_does_not_show_summary( assert main(["checkout"]) == 0 out, _ = capsys.readouterr() - assert escape_ansi(out).splitlines() == [f"D\tdir{os.sep}", "D\tother"] + assert out.splitlines() == [ + f"D\tdir{os.sep}".expandtabs(), + "D\tother".expandtabs(), + ] def test_stats_does_not_show_changes_by_default(tmp_dir, dvc, scm, capsys): diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index dfd0725ab..2f204c356 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -18,7 +18,6 @@ from dvc.remote.base import ( from dvc.stage.exceptions import StageNotFound from dvc.utils.fs import move, remove from dvc.utils.serialize import dump_yaml, load_yaml -from tests.func.parsing.test_errors import escape_ansi all_clouds = [ pytest.lazy_fixture(cloud) @@ -492,9 +491,8 @@ def test_pull_stats(tmp_dir, dvc, capsys, caplog, local_remote): assert main(["pull", "--force"]) == 0 out, _ = capsys.readouterr() - out = escape_ansi(out) - assert "M\tbar" in out - assert "A\tfoo" in out + assert "M\tbar".expandtabs() in out + assert "A\tfoo".expandtabs() in out assert "2 files fetched" in caplog.text assert "1 file added" in caplog.text assert "1 file modified" in caplog.text diff --git a/tests/unit/command/test_checkout.py b/tests/unit/command/test_checkout.py index 6e76410e2..95014cb8e 100644 --- a/tests/unit/command/test_checkout.py +++ b/tests/unit/command/test_checkout.py @@ -1,6 +1,5 @@ from dvc.cli import parse_args from dvc.command.checkout import CmdCheckout, log_changes -from tests.func.parsing.test_errors import escape_ansi def test_checkout(tmp_dir, dvc, mocker): @@ -34,9 +33,9 @@ def test_log_changes(capsys): def _assert_output(stats, expected_outs): log_changes(stats) out, _ = capsys.readouterr() - actual_output = escape_ansi(out).splitlines() + actual_output = out.splitlines() for out, line in zip_longest(expected_outs, actual_output): - assert out in line + assert out.expandtabs() in line _assert_output(stats, ["M\tfile2", "A\tfile1", "A\tdir1/", "D\tdir2/"])