ui: Use sys.stdout and sys.stderr dynamically (#5729)

Before, we used to set it during initialization. But this
made it hard to test using capsys as capsys changes the
reference to sys.stdout and sys.stderr with the mocked one
during the test. But since we are using dvc.ui globally and
setting it to a real sys.stdout and sys.stderr before capsys
comes into effect, testing was not possible with it.

Now, by using property hack, we will be using the values of
sys.stdout and sys.stderr during the time of write, so the mocking
would work and would be easier to test.

dvc/ui/__init__.py

@@ -29,13 +29,19 @@ class Console:
         error: TextIO = None,
         enable: bool = False,
     ) -> None:
-        self._input: TextIO = sys.stdin
-        self._output: TextIO = output or sys.stdout
-        self._error: TextIO = error or sys.stderr
-
+        self._output: Optional[TextIO] = output
+        self._error: Optional[TextIO] = error
         self.formatter: Formatter = formatter or Formatter()
         self._enabled: bool = enable
 
+    @property
+    def output(self) -> TextIO:
+        return self._output or sys.stdout
+
+    @property
+    def error_output(self) -> TextIO:
+        return self._error or sys.stderr
+
     def enable(self):
         self._enabled = True
 
@@ -61,7 +67,7 @@ class Console:
             style=style,
             sep=sep,
             end=end,
-            file=self._error,
+            file=self.error_output,
             flush=flush,
         )
 
@@ -77,12 +83,12 @@ class Console:
         if not self._enabled:
             return
 
-        file = file or self._output
+        file = file or self.output
         values = (self.formatter.format(obj, style=style) for obj in objects)
         return print(*values, sep=sep, end=end, file=file, flush=flush)
 
     def progress(self, *args, **kwargs) -> Tqdm:
-        kwargs.setdefault("file", self._error)
+        kwargs.setdefault("file", self.error_output)
         return Tqdm(*args, **kwargs)
 
     def prompt(
@@ -125,7 +131,7 @@ class Console:
         """rich_console is only set to stdout for now."""
         from rich import console
 
-        return console.Console(file=self._output)
+        return console.Console(file=self.output)
 
     def rich_table(self, pager: bool = True):
         pass

tests/conftest.py

@@ -43,6 +43,13 @@ def reset_loglevel(request, caplog):
         yield
 
 
+@pytest.fixture(autouse=True)
+def enable_ui():
+    from dvc.ui import ui
+
+    ui.enable()
+
+
 @pytest.fixture(scope="session", autouse=True)
 def _close_pools():
     from dvc.fs.pool import close_pools

tests/unit/ui/test_console.py

@@ -0,0 +1,28 @@
+from pytest import CaptureFixture
+
+from dvc.ui import Console
+
+
+def test_write(capsys: CaptureFixture[str]):
+    """Test that ui.write works."""
+    console = Console(enable=True)
+    message = "hello world"
+    console.write(message)
+    console.error_write(message)
+
+    captured = capsys.readouterr()
+    assert captured.out == f"{message}\n"
+    assert captured.err == f"{message}\n"
+
+
+def test_capsys_works(capsys: CaptureFixture[str]):
+    """Sanity check that capsys can capture outputs from a global ui."""
+    from dvc.ui import ui
+
+    message = "hello world"
+    ui.write(message)
+    ui.error_write(message)
+
+    captured = capsys.readouterr()
+    assert captured.out == f"{message}\n"
+    assert captured.err == f"{message}\n"