Merge pull request #2038 from Textualize/print-json-indent-fix

Fix highlighting issue when printing JSON that isn't indented

Author: willmcgugan

CHANGELOG.md

@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Collapsed definitions for single-character spinners, to save memory and reduce import time.
 - Fix print_json indent type in __init__.py
 - Fix error when inspecting object defined in REPL https://github.com/Textualize/rich/pull/2037
+- Fix incorrect highlighting of non-indented JSON https://github.com/Textualize/rich/pull/2038
 
 ### Changed
 

rich/console.py

@@ -2209,3 +2209,5 @@ def save_html(
         }
     )
     console.log("foo")
+
+    console.print_json(data={"name": "apple", "count": 1}, indent=None)

rich/highlighter.py

@@ -1,7 +1,9 @@
+import re
+import string
 from abc import ABC, abstractmethod
 from typing import List, Union
 
-from .text import Text
+from .text import Span, Text
 
 
 def _combine_regex(*regexes: str) -> str:
@@ -104,17 +106,39 @@ class ReprHighlighter(RegexHighlighter):
 class JSONHighlighter(RegexHighlighter):
     """Highlights JSON"""
 
+    # Captures the start and end of JSON strings, handling escaped quotes
+    JSON_STR = r"(?<![\\\w])(?P<str>b?\".*?(?<!\\)\")"
+    JSON_WHITESPACE = {" ", "\n", "\r", "\t"}
+
     base_style = "json."
     highlights = [
         _combine_regex(
             r"(?P<brace>[\{\[\(\)\]\}])",
             r"\b(?P<bool_true>true)\b|\b(?P<bool_false>false)\b|\b(?P<null>null)\b",
             r"(?P<number>(?<!\w)\-?[0-9]+\.?[0-9]*(e[\-\+]?\d+?)?\b|0x[0-9a-fA-F]*)",
-            r"(?<![\\\w])(?P<str>b?\".*?(?<!\\)\")",
+            JSON_STR,
         ),
-        r"(?<![\\\w])(?P<key>b?\".*?(?<!\\)\")\:",
     ]
 
+    def highlight(self, text: Text) -> None:
+        super().highlight(text)
+
+        # Additional work to handle highlighting JSON keys
+        plain = text.plain
+        append = text.spans.append
+        whitespace = self.JSON_WHITESPACE
+        for match in re.finditer(self.JSON_STR, plain):
+            start, end = match.span()
+            cursor = end
+            while cursor < len(plain):
+                char = plain[cursor]
+                cursor += 1
+                if char == ":":
+                    append(Span(start, end, "json.key"))
+                elif char in whitespace:
+                    continue
+                break
+
 
 if __name__ == "__main__":  # pragma: no cover
     from .console import Console
@@ -145,3 +169,6 @@ class JSONHighlighter(RegexHighlighter):
     console.print(
         "127.0.1.1 bar 192.168.1.4 2001:0db8:85a3:0000:0000:8a2e:0370:7334 foo"
     )
+    import json
+
+    console.print_json(json.dumps(obj={"name": "apple", "count": 1}), indent=None)

tests/test_console.py

@@ -14,8 +14,8 @@
     Console,
     ConsoleDimensions,
     ConsoleOptions,
-    group,
     ScreenUpdate,
+    group,
 )
 from rich.control import Control
 from rich.measure import measure_renderables
@@ -185,6 +185,15 @@ def test_print_json_ensure_ascii():
     assert result == expected
 
 
+def test_print_json_indent_none():
+    console = Console(file=io.StringIO(), color_system="truecolor")
+    data = {"name": "apple", "count": 1}
+    console.print_json(data=data, indent=None)
+    result = console.file.getvalue()
+    expected = '\x1b[1m{\x1b[0m\x1b[1;34m"name"\x1b[0m: \x1b[32m"apple"\x1b[0m, \x1b[1;34m"count"\x1b[0m: \x1b[1;36m1\x1b[0m\x1b[1m}\x1b[0m\n'
+    assert result == expected
+
+
 def test_log():
     console = Console(
         file=io.StringIO(),

tests/test_highlighter.py

@@ -1,8 +1,10 @@
 """Tests for the highlighter classes."""
-import pytest
+import json
 from typing import List
 
-from rich.highlighter import NullHighlighter, ReprHighlighter
+import pytest
+
+from rich.highlighter import JSONHighlighter, NullHighlighter, ReprHighlighter
 from rich.text import Span, Text
 
 
@@ -92,3 +94,53 @@ def test_highlight_regex(test: str, spans: List[Span]):
     highlighter.highlight(text)
     print(text.spans)
     assert text.spans == spans
+
+
+def test_highlight_json_with_indent():
+    json_string = json.dumps({"name": "apple", "count": 1}, indent=4)
+    text = Text(json_string)
+    highlighter = JSONHighlighter()
+    highlighter.highlight(text)
+    assert text.spans == [
+        Span(0, 1, "json.brace"),
+        Span(6, 12, "json.str"),
+        Span(14, 21, "json.str"),
+        Span(27, 34, "json.str"),
+        Span(36, 37, "json.number"),
+        Span(38, 39, "json.brace"),
+        Span(6, 12, "json.key"),
+        Span(27, 34, "json.key"),
+    ]
+
+
+def test_highlight_json_string_only():
+    json_string = '"abc"'
+    text = Text(json_string)
+    highlighter = JSONHighlighter()
+    highlighter.highlight(text)
+    assert text.spans == [Span(0, 5, "json.str")]
+
+
+def test_highlight_json_empty_string_only():
+    json_string = '""'
+    text = Text(json_string)
+    highlighter = JSONHighlighter()
+    highlighter.highlight(text)
+    assert text.spans == [Span(0, 2, "json.str")]
+
+
+def test_highlight_json_no_indent():
+    json_string = json.dumps({"name": "apple", "count": 1}, indent=None)
+    text = Text(json_string)
+    highlighter = JSONHighlighter()
+    highlighter.highlight(text)
+    assert text.spans == [
+        Span(0, 1, "json.brace"),
+        Span(1, 7, "json.str"),
+        Span(9, 16, "json.str"),
+        Span(18, 25, "json.str"),
+        Span(27, 28, "json.number"),
+        Span(28, 29, "json.brace"),
+        Span(1, 7, "json.key"),
+        Span(18, 25, "json.key"),
+    ]