Fix JSON key highlighting when no indentation

rich/highlighter.py

@@ -1,7 +1,10 @@
+import json
+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 +107,38 @@ 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?\".*?(?<!\\)\")"
+
     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
+        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 in string.whitespace:
+                    continue
+                elif char == ":":
+                    append(Span(start, end, "json.key"))
+                else:
+                    break
+
 
 if __name__ == "__main__":  # pragma: no cover
     from .console import Console
@@ -145,3 +169,5 @@ class JSONHighlighter(RegexHighlighter):
     console.print(
         "127.0.1.1 bar 192.168.1.4 2001:0db8:85a3:0000:0000:8a2e:0370:7334 foo"
     )
+
+    console.print_json(json.dumps(obj={"name": "apple", "count": 1}), indent=None)

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,37 @@ 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_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"),
+    ]