From 058377f33e50b733586c3fb9bdf15ad7edd3d021 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Mon, 27 Sep 2021 18:11:20 +0545 Subject: [PATCH] ui: support pager universally (#6692) --- dvc/command/dag.py | 5 ++--- dvc/ui/__init__.py | 32 ++++++++++++++++++++++++++------ dvc/{utils => ui}/pager.py | 0 dvc/ui/table.py | 28 ++++++++++++---------------- tests/unit/ui/test_pager.py | 4 ++-- tests/unit/ui/test_table.py | 7 ++++--- 6 files changed, 46 insertions(+), 30 deletions(-) rename dvc/{utils => ui}/pager.py (100%) diff --git a/dvc/command/dag.py b/dvc/command/dag.py index 8867bf1dd..f86f8605f 100644 --- a/dvc/command/dag.py +++ b/dvc/command/dag.py @@ -103,9 +103,8 @@ class CmdDAG(CmdBase): if self.args.dot: ui.write(_show_dot(G)) else: - from dvc.utils.pager import pager - - pager(_show_ascii(G)) + with ui.pager(): + ui.write(_show_ascii(G)) return 0 diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py index ff4898d2f..9b515f252 100644 --- a/dvc/ui/__init__.py +++ b/dvc/ui/__init__.py @@ -1,9 +1,10 @@ -import sys +from contextlib import contextmanager from typing import ( TYPE_CHECKING, Any, Dict, Iterable, + Iterator, Optional, Sequence, TextIO, @@ -42,8 +43,13 @@ class Console: def __init__( self, formatter: Formatter = None, enable: bool = False ) -> None: + from contextvars import ContextVar + self.formatter: Formatter = formatter or Formatter() self._enabled: bool = enable + self._paginate: ContextVar[bool] = ContextVar( + "_paginate", default=False + ) def enable(self) -> None: self._enabled = True @@ -87,6 +93,8 @@ class Console: styled: bool = False, file: TextIO = None, ) -> None: + import sys + from dvc.progress import Tqdm sep = " " if sep is None else sep @@ -96,13 +104,14 @@ class Console: file = file or (sys.stderr if stderr else sys.stdout) with Tqdm.external_write_mode(file=file): - if styled: + # if we are inside pager context, send the output to rich's buffer + if styled or self._paginate.get(): console = self.error_console if stderr else self.rich_console - return console.print(*objects, sep=sep, end=end) + if styled: + return console.print(*objects, sep=sep, end=end) + return console.out(*objects, sep=sep, end=end, highlight=False) - values = ( - self.formatter.format(obj, style=style) for obj in objects - ) + values = (self.formatter.format(obj, style) for obj in objects) return print(*values, sep=sep, end=end, file=file) @property @@ -117,6 +126,17 @@ class Console: return Tqdm(*args, **kwargs) + @contextmanager + def pager(self, styles: bool = True) -> Iterator[None]: + from .pager import DvcPager + + tok = self._paginate.set(True) + try: + with self.rich_console.pager(pager=DvcPager(), styles=styles): + yield + finally: + self._paginate.reset(tok) + def prompt( self, text: str, choices: Iterable[str] = None, password: bool = False ) -> Optional[str]: diff --git a/dvc/utils/pager.py b/dvc/ui/pager.py similarity index 100% rename from dvc/utils/pager.py rename to dvc/ui/pager.py diff --git a/dvc/ui/table.py b/dvc/ui/table.py index de795dcfc..ae17613df 100644 --- a/dvc/ui/table.py +++ b/dvc/ui/table.py @@ -1,5 +1,5 @@ from collections import abc -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from itertools import zip_longest from typing import TYPE_CHECKING, Dict, Iterator, Sequence, Union @@ -11,7 +11,6 @@ if TYPE_CHECKING: from dvc.ui import Console, RichText - SHOW_MAX_WIDTH = 1024 @@ -30,6 +29,7 @@ def plain_table( pager: bool = False, force: bool = True, ) -> None: + from funcy import nullcontext from tabulate import tabulate text: str = tabulate( @@ -44,11 +44,8 @@ def plain_table( # NOTE: md table is incomplete without the trailing newline text += "\n" - if pager: - from dvc.utils.pager import pager as _pager - - _pager(text) - else: + cm = ui.pager() if pager else nullcontext() + with cm: ui.write(text, force=force) @@ -110,14 +107,13 @@ def rich_table( for row, style in zip_longest(data, rs): table.add_row(*row, **(style or {})) - console = ui.rich_console + stack = ExitStack() + if pager: + stack.enter_context( + console_width(table, ui.rich_console, SHOW_MAX_WIDTH) + ) + stack.enter_context(ui.pager()) - if not pager: - console.print(table) + with stack: + ui.write(table, styled=True) 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/tests/unit/ui/test_pager.py b/tests/unit/ui/test_pager.py index eacc3dfd2..0d9984ca1 100644 --- a/tests/unit/ui/test_pager.py +++ b/tests/unit/ui/test_pager.py @@ -1,7 +1,7 @@ import pytest from dvc.env import DVC_PAGER -from dvc.utils.pager import ( +from dvc.ui.pager import ( DEFAULT_PAGER, LESS, PAGER_ENV, @@ -91,7 +91,7 @@ def test_pager(mocker, monkeypatch): monkeypatch.setenv(DVC_PAGER, "my-pager") mocker.patch("sys.stdout.isatty", return_value=True) - m_make_pager = mocker.patch("dvc.utils.pager.make_pager") + m_make_pager = mocker.patch("dvc.ui.pager.make_pager") _pager = m_make_pager.return_value = mocker.MagicMock() pager("hello world") diff --git a/tests/unit/ui/test_table.py b/tests/unit/ui/test_table.py index 96a2160cf..71126d855 100644 --- a/tests/unit/ui/test_table.py +++ b/tests/unit/ui/test_table.py @@ -43,7 +43,7 @@ def test_plain_md(capsys: CaptureFixture[str]): def test_plain_pager(mocker: MockerFixture): - pager_mock = mocker.patch("dvc.utils.pager.pager") + pager_mock = mocker.patch("dvc.ui.pager.pager") ui.table( [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")], headers=["first", "second"], @@ -56,7 +56,8 @@ def test_plain_pager(mocker: MockerFixture): first second foo bar foo1 bar1 - foo2 bar2""" + foo2 bar2 + """ ) ) @@ -138,7 +139,7 @@ def test_rich_styles(capsys: CaptureFixture[str], extra_opts): def test_rich_pager(mocker: MockerFixture): - pager_mock = mocker.patch("dvc.utils.pager.pager") + pager_mock = mocker.patch("dvc.ui.pager.pager") ui.table( [("foo", "bar"), ("foo1", "bar1"), ("foo2", "bar2")],