diff --git a/dvc/command/diff.py b/dvc/command/diff.py index b6023e364..85e4f5a75 100644 --- a/dvc/command/diff.py +++ b/dvc/command/diff.py @@ -1,5 +1,4 @@ import argparse -import json import logging import os @@ -143,7 +142,7 @@ class CmdDiff(CmdBase): diff[key] = entries if self.args.json: - ui.write(json.dumps(diff)) + ui.write_json(diff) elif self.args.markdown: _show_markdown(diff, show_hash, hide_missing) elif diff: diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index 6b3ce16da..f40adfc86 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -505,9 +505,7 @@ class CmdExperimentsShow(CmdBase): return 1 if self.args.json: - import json - - ui.write(json.dumps(all_experiments, default=_format_json)) + ui.write_json(all_experiments, default=_format_json) else: precision = ( self.args.precision or None @@ -560,9 +558,7 @@ class CmdExperimentsDiff(CmdBase): return 1 if self.args.json: - import json - - ui.write(json.dumps(diff)) + ui.write_json(diff) else: from dvc.compare import show_diff diff --git a/dvc/command/metrics.py b/dvc/command/metrics.py index 57261b0d8..d41bc7215 100644 --- a/dvc/command/metrics.py +++ b/dvc/command/metrics.py @@ -32,9 +32,7 @@ class CmdMetricsShow(CmdMetricsBase): return 1 if self.args.json: - import json - - ui.write(json.dumps(metrics, default=encode_exception)) + ui.write_json(metrics, default=encode_exception) else: from dvc.compare import show_metrics @@ -66,9 +64,7 @@ class CmdMetricsDiff(CmdMetricsBase): return 1 if self.args.json: - import json - - ui.write(json.dumps(diff)) + ui.write_json(diff) else: from dvc.compare import show_diff diff --git a/dvc/command/params.py b/dvc/command/params.py index 199427c0e..5796cb8c0 100644 --- a/dvc/command/params.py +++ b/dvc/command/params.py @@ -26,9 +26,7 @@ class CmdParamsDiff(CmdBase): return 1 if self.args.json: - import json - - ui.write(json.dumps(diff)) + ui.write_json(diff) else: from dvc.compare import show_diff diff --git a/dvc/command/status.py b/dvc/command/status.py index 1b74f096b..6754843c6 100644 --- a/dvc/command/status.py +++ b/dvc/command/status.py @@ -70,9 +70,7 @@ class CmdDataStatus(CmdDataBase): return bool(st) if self.args.json: - import json - - ui.write(json.dumps(st)) + ui.write_json(st) return 0 if st: diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py index 80ce1586d..f91e87c69 100644 --- a/dvc/ui/__init__.py +++ b/dvc/ui/__init__.py @@ -2,6 +2,7 @@ from contextlib import contextmanager from typing import ( TYPE_CHECKING, Any, + Callable, Dict, Iterable, Iterator, @@ -83,17 +84,20 @@ class Console: ) def write_json( - self, - data: Any, - indent: int = 2, + self, data: Any, indent: int = 2, default: Callable = None ) -> None: - if self.isatty(): - from rich.json import JSON - - j = JSON.from_data(data, indent=indent) - return self.write(j, styled=True) import json + if self.isatty(): + from rich.highlighter import JSONHighlighter + + j = json.dumps(data, indent=indent, default=default) + highlighter = JSONHighlighter() + text = highlighter(j) + text.no_wrap = True + text.overflow = None + return self.write(text, styled=True) + return self.write(json.dumps(data)) def write( diff --git a/tests/unit/ui/test_console.py b/tests/unit/ui/test_console.py index 9e3680493..237e4bca9 100644 --- a/tests/unit/ui/test_console.py +++ b/tests/unit/ui/test_console.py @@ -1,4 +1,8 @@ +import textwrap + +import pytest from pytest import CaptureFixture +from pytest_mock import MockerFixture from dvc.ui import Console @@ -15,6 +19,42 @@ def test_write(capsys: CaptureFixture[str]): assert captured.err == f"{message}\n" +@pytest.mark.parametrize( + "isatty, expected_output", + [ + ( + True, + textwrap.dedent( + """\ + { + "hello": "world" + } + """ + ), + ), + ( + False, + textwrap.dedent( + """\ + {"hello": "world"} + """ + ), + ), + ], +) +def test_write_json( + capsys: CaptureFixture[str], mocker: MockerFixture, isatty, expected_output +): + """Test that ui.write_json works.""" + + console = Console(enable=True) + mocker.patch.object(console, "isatty", return_value=isatty) + message = {"hello": "world"} + console.write_json(message) + captured = capsys.readouterr() + assert captured.out == expected_output + + def test_capsys_works(capsys: CaptureFixture[str]): """Sanity check that capsys can capture outputs from a global ui.""" from dvc.ui import ui