diff --git a/dvc/types.py b/dvc/types.py index 837e5cd52..94aceddf6 100644 --- a/dvc/types.py +++ b/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] diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py index 164a859ad..737a0778b 100644 --- a/dvc/ui/__init__.py +++ b/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 ) diff --git a/dvc/ui/table.py b/dvc/ui/table.py new file mode 100644 index 000000000..aaf6d10dd --- /dev/null +++ b/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) diff --git a/dvc/utils/pager.py b/dvc/utils/pager.py index bf3303159..f758b351d 100644 --- a/dvc/utils/pager.py +++ b/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) diff --git a/tests/unit/ui/test_table.py b/tests/unit/ui/test_table.py new file mode 100644 index 000000000..663175999 --- /dev/null +++ b/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", "")