command: output prettify json in tty (#6743)

1. update diff.py,experiments.py,metrics.py,params.py,status.py to
print prettify json
2. rewrite write_json in __init__.py to accept default function
3. add new unit test for write_json in tests/unit/ui/test_console.py
4. Fix #6712

Co-authored-by: Zhenyu Gao <[email protected]>

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:

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
 

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
 

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
 

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:

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(

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