diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py index 1b6cd6c8b..164a859ad 100644 --- a/dvc/ui/__init__.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index e5d906fa1..4a000c5b4 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/unit/ui/__init__.py b/tests/unit/ui/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/ui/test_console.py b/tests/unit/ui/test_console.py new file mode 100644 index 000000000..9e3680493 --- /dev/null +++ b/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"