ui: unify tables (#5756)

Author: skshetry

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 @@
 
 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 @@ def output(self) -> TextIO:
     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 @@ def write(
         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 @@ def confirm(self, statement: str) -> bool:
             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, header, rows, markdown: bool = False):
-        from tabulate import tabulate
+    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
 
-        if not rows and not markdown:
-            return ""
+        if not data and not markdown:
+            return
 
-        ret = tabulate(
-            rows,
-            header,
-            tablefmt="github" if markdown else "plain",
-            disable_numparse=True,
-            # None will be shown as "" by default, overriding
-            missingval="—",
-        )
+        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,
+                )
 
-        if markdown:
-            # NOTE: md table is incomplete without the trailing newline
-            ret += "\n"
+            return
 
-        self.write(ret)
+        return t.plain_table(
+            self, data, headers, markdown=markdown, pager=pager, force=force,
+        )
 
 
 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)