ui: unify tables (#5756)

dvc/types.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from dvc.path_info import PathInfo, URLInfo
 
@@ -11,3 +11,5 @@ AnyPath = Union[str, DvcPath, StrPath]
 
 OptStr = Optional[str]
 TargetType = Union[List[str], str]
+DictStrAny = Dict[str, Any]
+DictAny = Dict[Any, Any]

dvc/ui/__init__.py

@@ -1,12 +1,22 @@
 import sys
 from collections import defaultdict
-from typing import Any, Dict, Iterable, Optional, TextIO
-
-from funcy import cached_property
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Optional,
+    Sequence,
+    TextIO,
+    Union,
+)
 
 from dvc.progress import Tqdm
 from dvc.utils import colorize
 
+if TYPE_CHECKING:
+    from dvc.ui.table import Headers, Styles, TableData
+
 
 class Formatter:
     def __init__(self, theme: Dict = None, defaults: Dict = None) -> None:
@@ -42,7 +52,7 @@ class Console:
     def error_output(self) -> TextIO:
         return self._error or sys.stderr
 
-    def enable(self):
+    def enable(self) -> None:
         self._enabled = True
 
     def success(self, message: str) -> None:
@@ -79,8 +89,9 @@ class Console:
         end: str = None,
         file: TextIO = None,
         flush: bool = False,
+        force: bool = False,
     ) -> None:
-        if not self._enabled:
+        if not self._enabled and not force:
             return
 
         file = file or self.output
@@ -126,49 +137,64 @@ class Console:
             return False
         return answer.startswith("y")
 
-    @cached_property
+    @property
     def rich_console(self):
         """rich_console is only set to stdout for now."""
         from rich import console
 
+        # FIXME: Getting IO Operation on closed file error
+        #  when testing with capsys, therefore we are creating
+        #  one instance each time as a temporary workaround.
         return console.Console(file=self.output)
 
-    def rich_table(self, pager: bool = True):
-        pass
+    def table(
+        self,
+        data: "TableData",
+        headers: "Headers" = None,
+        markdown: bool = False,
+        rich_table: bool = False,
+        force: bool = True,
+        pager: bool = False,
+        header_styles: Sequence["Styles"] = None,
+        row_styles: Sequence["Styles"] = None,
+        borders: Union[bool, str] = False,
+    ) -> None:
+        from dvc.ui import table as t
 
-    def table(self, header, rows, markdown: bool = False):
-        from tabulate import tabulate
+        if not data and not markdown:
+            return
 
-        if not rows and not markdown:
-            return ""
+        if not markdown and rich_table:
+            if force or self._enabled:
+                return t.rich_table(
+                    self,
+                    data,
+                    headers,
+                    pager=pager,
+                    header_styles=header_styles,
+                    row_styles=row_styles,
+                    borders=borders,
+                )
 
-        ret = tabulate(
-            rows,
-            header,
-            tablefmt="github" if markdown else "plain",
-            disable_numparse=True,
-            # None will be shown as "" by default, overriding
-            missingval="—",
+            return
+
+        return t.plain_table(
+            self, data, headers, markdown=markdown, pager=pager, force=force,
         )
 
-        if markdown:
-            # NOTE: md table is incomplete without the trailing newline
-            ret += "\n"
-
-        self.write(ret)
-
 
 ui = Console()
 
 
 if __name__ == "__main__":
     ui.enable()
+
     ui.write("No default remote set")
     ui.success("Everything is up to date.")
     ui.warn("Run queued experiments will be removed.")
     ui.error("too few arguments.")
 
-    ui.table("keys", {"Path": ["scores.json"], "auc": ["0.5674"]})
+    ui.table([("scores.json", "0.5674")], headers=["Path", "auc"])
     ui.table(
-        "keys", {"Path": ["scores.json"], "auc": ["0.5674"]}, markdown=True
+        [("scores.json", "0.5674")], headers=["Path", "auc"], markdown=True
     )

dvc/ui/table.py

@@ -0,0 +1,119 @@
+from contextlib import contextmanager
+from itertools import zip_longest
+from typing import TYPE_CHECKING, Iterator, List, Sequence, Union
+
+from dvc.types import DictStrAny
+
+if TYPE_CHECKING:
+    from rich.console import Console as RichConsole
+    from rich.table import Table
+    from rich.text import Text
+
+    from dvc.ui import Console
+
+
+SHOW_MAX_WIDTH = 1024
+
+
+CellT = Union[str, "Text"]  # Text is mostly compatible with str
+Row = Sequence[CellT]
+TableData = List[Row]
+Headers = Sequence[str]
+Styles = DictStrAny
+
+
+def plain_table(
+    ui: "Console",
+    data: TableData,
+    headers: Headers = None,
+    markdown: bool = False,
+    pager: bool = False,
+    force: bool = True,
+) -> None:
+    from tabulate import tabulate
+
+    text: str = tabulate(
+        data,
+        headers if headers is not None else (),
+        tablefmt="github" if markdown else "plain",
+        disable_numparse=True,
+        # None will be shown as "" by default, overriding
+        missingval="-",
+    )
+    if markdown:
+        # NOTE: md table is incomplete without the trailing newline
+        text += "\n"
+
+    if pager:
+        from dvc.utils.pager import pager as _pager
+
+        _pager(text)
+    else:
+        ui.write(text, force=force)
+
+
+@contextmanager
+def console_width(
+    table: "Table", console: "RichConsole", val: int
+) -> Iterator[None]:
+    # NOTE: rich does not have native support for unlimited width
+    # via pager. we override rich table compression by setting
+    # console width to the full width of the table
+    # pylint: disable=protected-access
+
+    console_options = console.options
+    original = console_options.max_width
+    con_width = console._width
+
+    try:
+        console_options.max_width = val
+        measurement = table.__rich_measure__(console, console_options)
+        console._width = measurement.maximum
+
+        yield
+    finally:
+        console_options.max_width = original
+        console._width = con_width
+
+
+def rich_table(
+    ui: "Console",
+    data: TableData,
+    headers: Headers = None,
+    pager: bool = False,
+    header_styles: Sequence[Styles] = None,
+    row_styles: Sequence[Styles] = None,
+    borders: Union[bool, str] = False,
+) -> None:
+    from rich import box
+
+    from dvc.utils.table import Table
+
+    border_style = {
+        True: box.HEAVY_HEAD,  # is a default in rich,
+        False: None,
+        "simple": box.SIMPLE,
+        "minimal": box.MINIMAL,
+    }
+
+    table = Table(box=border_style[borders])
+    hs: Sequence[Styles] = header_styles or []
+    rs: Sequence[Styles] = row_styles or []
+
+    for header, style in zip_longest(headers or [], hs):
+        table.add_column(header, **(style or {}))
+
+    for row, style in zip_longest(data, rs):
+        table.add_row(*row, **(style or {}))
+
+    console = ui.rich_console
+
+    if not pager:
+        console.print(table)
+        return
+
+    from dvc.utils.pager import DvcPager
+
+    with console_width(table, console, SHOW_MAX_WIDTH):
+        with console.pager(pager=DvcPager(), styles=True):
+            console.print(table)

dvc/utils/pager.py

@@ -46,7 +46,7 @@ def find_pager():
     return pydoc.plainpager
 
 
-def pager(text):
+def pager(text: str) -> None:
     find_pager()(text)
 
 

tests/unit/ui/test_table.py

@@ -0,0 +1,164 @@
+import textwrap
+
+import pytest
+from pytest import CaptureFixture
+from pytest_mock import MockerFixture
+from rich.style import Style
+
+from dvc.ui import ui
+
+
+def test_plain(capsys: CaptureFixture[str]):
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+    )
+    captured = capsys.readouterr()
+    assert captured.out == textwrap.dedent(
+        """\
+        first    second
+        foo      bar
+        foo1     bar1
+        foo2     bar2
+    """
+    )
+
+
+def test_plain_md(capsys: CaptureFixture[str]):
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+        markdown=True,
+    )
+    captured = capsys.readouterr()
+    assert captured.out == textwrap.dedent(
+        """\
+        | first   | second   |
+        |---------|----------|
+        | foo     | bar      |
+        | foo1    | bar1     |
+        | foo2    | bar2     |\n
+    """
+    )
+
+
+def test_plain_pager(mocker: MockerFixture):
+    pager_mock = mocker.patch("dvc.utils.pager.pager")
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+        pager=True,
+    )
+
+    pager_mock.assert_called_once_with(
+        textwrap.dedent(
+            """\
+            first    second
+            foo      bar
+            foo1     bar1
+            foo2     bar2"""
+        )
+    )
+
+
+def test_plain_headerless(capsys: CaptureFixture[str]):
+    ui.table([("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],)
+    captured = capsys.readouterr()
+    assert captured.out == textwrap.dedent(
+        """\
+        foo   bar
+        foo1  bar1
+        foo2  bar2
+    """
+    )
+
+
+def test_rich_simple(capsys: CaptureFixture[str]):
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+        rich_table=True,
+    )
+    # not able to test the actual style for now
+    captured = capsys.readouterr()
+    assert [
+        row.strip() for row in captured.out.splitlines() if row.strip()
+    ] == ["first  second", "foo    bar", "foo1   bar1", "foo2   bar2"]
+
+
+def test_rich_headerless(capsys: CaptureFixture[str]):
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")], rich_table=True,
+    )
+    captured = capsys.readouterr()
+    assert [
+        row.strip() for row in captured.out.splitlines() if row.strip()
+    ] == ["foo   bar", "foo1  bar1", "foo2  bar2"]
+
+
+def test_rich_border(capsys: CaptureFixture[str]):
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+        rich_table=True,
+        borders="simple",
+    )
+    captured = capsys.readouterr()
+    assert [
+        row.strip() for row in captured.out.splitlines() if row.strip()
+    ] == [
+        "first   second",
+        "────────────────",
+        "foo     bar",
+        "foo1    bar1",
+        "foo2    bar2",
+    ]
+
+
+@pytest.mark.parametrize(
+    "extra_opts",
+    [
+        {"header_styles": [{"style": Style(bold=True)}]},
+        {"row_styles": [{"style": Style(bold=True)}]},
+    ],
+)
+def test_rich_styles(capsys: CaptureFixture[str], extra_opts):
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+        rich_table=True,
+        **extra_opts
+    )
+    # not able to test the actual style for now
+    captured = capsys.readouterr()
+    assert [
+        row.strip() for row in captured.out.splitlines() if row.strip()
+    ] == ["first  second", "foo    bar", "foo1   bar1", "foo2   bar2"]
+
+
+def test_rich_pager(mocker: MockerFixture):
+    pager_mock = mocker.patch("dvc.utils.pager.pager")
+
+    ui.table(
+        [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],
+        headers=["first", "second"],
+        rich_table=True,
+        pager=True,
+    )
+    received_text = pager_mock.call_args[0][0]
+    assert [
+        row.strip() for row in received_text.splitlines() if row.strip()
+    ] == ["first  second", "foo    bar", "foo1   bar1", "foo2   bar2"]
+
+
+@pytest.mark.parametrize("rich_table", [True, False])
+def test_empty(capsys: CaptureFixture[str], rich_table: str):
+    ui.table([], rich_table=rich_table)
+    out, err = capsys.readouterr()
+    assert (out, err) == ("", "")
+
+
+def test_empty_markdown(capsys: CaptureFixture[str]):
+    ui.table([], headers=["Col1", "Col2"], markdown=True)
+    out, err = capsys.readouterr()
+    assert (out, err) == ("| Col1   | Col2   |\n|--------|--------|\n\n", "")