ui: write gets a styled param for rich console printing (#6037)

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):

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",

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):

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

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/"])