ui: support pager universally (#6692)

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
 

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

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)

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

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