improve json double quote

pretty_json.py

from __future__ import annotations

import argparse
import io
import json
import re
import sys
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

from rich.console import Console

type JSONValue = dict[str, "JSONValue"] | list["JSONValue"] | str | int | float | bool | None
type JSONPathPart = str | int
type JSONPath = tuple[JSONPathPart, ...]

DEFAULT_MAX_DECODE_DEPTH = 20
DECODED_STRING_COMMENT = "/* jsonx: was a JSON string */"
INLINE_STRING_COMMENT = "/* jsonx: was a string containing inline JSON */"
INLINE_VALUE_COMMENT = "/* jsonx: extracted from string */"
JSONX_COMMENT_HIGHLIGHT = "\x1b[1;30;103m"
JSON_KEY_COLOR = "\x1b[36m"
JSON_STRING_COLOR = "\x1b[32m"
JSON_LITERAL_COLOR = "\x1b[35m"
JSON_PUNCT_COLOR = "\x1b[90m"
ANSI_RESET = "\x1b[0m"
JSON_NUMBER_PATTERN = re.compile(r"-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?")


@dataclass(slots=True)
class ConvertedDocument:
    value: JSONValue
    decoded_string_paths: set[JSONPath]


@dataclass(slots=True)
class InlineJsonSegment:
    label: str
    start: int
    end: int
    value: JSONValue


def convert_json_strings(value: JSONValue, *, max_decode_depth: int = DEFAULT_MAX_DECODE_DEPTH) -> JSONValue:
    converted, _ = _convert_json_strings_with_annotations(
        value=value,
        path=(),
        max_decode_depth=max_decode_depth,
    )
    return converted


def convert_text(text: str, *, max_decode_depth: int = DEFAULT_MAX_DECODE_DEPTH) -> str:
    return format_documents(
        documents=convert_documents(text=text, max_decode_depth=max_decode_depth),
        color=False,
    )


def convert_documents(text: str, *, max_decode_depth: int = DEFAULT_MAX_DECODE_DEPTH) -> list[JSONValue]:
    return [
        document.value for document in convert_documents_with_annotations(text=text, max_decode_depth=max_decode_depth)
    ]


def convert_documents_with_annotations(
    text: str, *, max_decode_depth: int = DEFAULT_MAX_DECODE_DEPTH
) -> list[ConvertedDocument]:
    documents = parse_json_documents(text=text)
    return [
        ConvertedDocument(value=converted, decoded_string_paths=decoded_string_paths)
        for document in documents
        for converted, decoded_string_paths in [
            _convert_json_strings_with_annotations(
                value=document,
                path=(),
                max_decode_depth=max_decode_depth,
            )
        ]
    ]


def format_documents(*, documents: list[JSONValue], color: bool) -> str:
    if not color:
        return "\n\n".join(json.dumps(document, ensure_ascii=False, indent=2) for document in documents) + "\n"

    output = io.StringIO()
    console = Console(file=output, force_terminal=True, color_system="standard")
    for index, document in enumerate(documents):
        if index > 0:
            console.print()
        console.print_json(json=json.dumps(document, ensure_ascii=False))
    return output.getvalue()


def format_converted_documents(*, documents: list[ConvertedDocument], color: bool, comments: bool) -> str:
    if not comments:
        return format_documents(documents=[document.value for document in documents], color=color)

    output_text = (
        "\n\n".join(
            _format_jsonc_value(
                value=document.value,
                decoded_string_paths=document.decoded_string_paths,
                path=(),
                indent=0,
            )
            for document in documents
        )
        + "\n"
    )
    if not color:
        return output_text

    return _highlight_jsonc_syntax(text=output_text)


def parse_json_documents(text: str) -> list[JSONValue]:
    decoder = json.JSONDecoder()
    documents: list[JSONValue] = []
    index = 0

    while True:
        index = _skip_whitespace(text=text, index=index)
        if index >= len(text):
            break
        try:
            document, index = decoder.raw_decode(text, index)
        except json.JSONDecodeError as exc:
            raise ValueError(f"invalid JSON at line {exc.lineno}, column {exc.colno}: {exc.msg}") from exc
        documents.append(document)

    if not documents:
        raise ValueError("input contains no JSON documents")
    return documents


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Pretty-print JSON and recursively expand stringified JSON values.")
    parser.add_argument(
        "input", nargs="?", default="-", help="Input JSON file. Defaults to stdin. Use '-' to read stdin."
    )
    parser.add_argument("-o", "--output", help="Write output to this file instead of stdout.")
    comment_group = parser.add_mutually_exclusive_group()
    comment_group.add_argument(
        "--comments",
        dest="comments",
        action="store_true",
        default=True,
        help="Emit JSONC comments marking converted values. This is the default.",
    )
    comment_group.add_argument(
        "--no-comments",
        dest="comments",
        action="store_false",
        help="Emit strict JSON without JSONC annotations.",
    )
    color_group = parser.add_mutually_exclusive_group()
    color_group.add_argument(
        "--color", dest="color", action="store_true", default=None, help="Force ANSI color output."
    )
    color_group.add_argument("--no-color", dest="color", action="store_false", help="Disable ANSI color output.")
    args = parser.parse_args(argv)

    try:
        input_text = sys.stdin.read() if args.input == "-" else Path(args.input).read_text(encoding="utf-8")

        color = args.color if args.color is not None else args.output is None and sys.stdout.isatty()
        output_text = format_converted_documents(
            documents=convert_documents_with_annotations(text=input_text),
            color=color,
            comments=args.comments,
        )

        if args.output:
            Path(args.output).write_text(output_text, encoding="utf-8")
        else:
            sys.stdout.write(output_text)
    except (OSError, ValueError) as exc:
        parser.exit(status=1, message=f"json-pretty-convert: {exc}\n")

    return 0


def _convert_json_strings_with_annotations(
    value: JSONValue,
    *,
    path: JSONPath,
    max_decode_depth: int,
) -> tuple[JSONValue, set[JSONPath]]:
    match value:
        case dict():
            converted: dict[str, JSONValue] = {}
            decoded_string_paths: set[JSONPath] = set()
            for key, child in value.items():
                child_value, child_paths = _convert_json_strings_with_annotations(
                    value=child,
                    path=(*path, key),
                    max_decode_depth=max_decode_depth,
                )
                converted[key] = child_value
                decoded_string_paths.update(child_paths)
            return converted, decoded_string_paths
        case list():
            converted_items: list[JSONValue] = []
            decoded_string_paths = set()
            for index, child in enumerate(value):
                child_value, child_paths = _convert_json_strings_with_annotations(
                    value=child,
                    path=(*path, index),
                    max_decode_depth=max_decode_depth,
                )
                converted_items.append(child_value)
                decoded_string_paths.update(child_paths)
            return converted_items, decoded_string_paths
        case str():
            return _decode_json_string_with_annotations(
                text=value,
                path=path,
                max_decode_depth=max_decode_depth,
            )
        case _:
            return value, set()


def _decode_json_string_with_annotations(
    *,
    text: str,
    path: JSONPath,
    max_decode_depth: int,
) -> tuple[JSONValue, set[JSONPath]]:
    current = text
    for _ in range(max_decode_depth):
        candidate = current.strip()
        if not _looks_like_embedded_json(text=candidate):
            return current, set()

        try:
            decoded = json.loads(candidate)
        except json.JSONDecodeError:
            return current, set()

        if isinstance(decoded, dict | list):
            converted, decoded_string_paths = _convert_json_strings_with_annotations(
                value=decoded,
                path=path,
                max_decode_depth=max_decode_depth,
            )
            decoded_string_paths.add(path)
            return converted, decoded_string_paths
        if isinstance(decoded, str):
            current = decoded
            continue
        return current, set()

    raise ValueError(f"exceeded nested JSON string decode depth at {_format_path(path=path)}")


def _format_jsonc_value(
    *,
    value: JSONValue,
    decoded_string_paths: set[JSONPath],
    path: JSONPath,
    indent: int,
) -> str:
    comment = f"{DECODED_STRING_COMMENT} " if path in decoded_string_paths else ""
    match value:
        case dict():
            body = _format_jsonc_object(
                value=value,
                decoded_string_paths=decoded_string_paths,
                path=path,
                indent=indent,
            )
        case list():
            body = _format_jsonc_array(
                value=value,
                decoded_string_paths=decoded_string_paths,
                path=path,
                indent=indent,
            )
        case str():
            body = _format_jsonc_string(value=value, path=path, indent=indent)
        case _:
            body = json.dumps(value, ensure_ascii=False)
    return f"{comment}{body}"


def _format_jsonc_object(
    *,
    value: dict[str, JSONValue],
    decoded_string_paths: set[JSONPath],
    path: JSONPath,
    indent: int,
) -> str:
    if not value:
        return "{}"

    child_indent = " " * (indent + 2)
    current_indent = " " * indent
    lines = ["{"]
    items = list(value.items())
    for index, (key, child) in enumerate(items):
        comma = "," if index < len(items) - 1 else ""
        child_path = (*path, key)
        child_text = _format_jsonc_value(
            value=child,
            decoded_string_paths=decoded_string_paths,
            path=child_path,
            indent=indent + 2,
        )
        lines.append(f"{child_indent}{json.dumps(key, ensure_ascii=False)}: {child_text}{comma}")
    lines.append(f"{current_indent}}}")
    return "\n".join(lines)


def _format_jsonc_array(
    *,
    value: list[JSONValue],
    decoded_string_paths: set[JSONPath],
    path: JSONPath,
    indent: int,
) -> str:
    if not value:
        return "[]"

    child_indent = " " * (indent + 2)
    current_indent = " " * indent
    lines = ["["]
    for index, child in enumerate(value):
        comma = "," if index < len(value) - 1 else ""
        child_text = _format_jsonc_value(
            value=child,
            decoded_string_paths=decoded_string_paths,
            path=(*path, index),
            indent=indent + 2,
        )
        lines.append(f"{child_indent}{child_text}{comma}")
    lines.append(f"{current_indent}]")
    return "\n".join(lines)


def _format_jsonc_string(*, value: str, path: JSONPath, indent: int) -> str:
    segments = _extract_inline_json_segments(text=value)
    if not segments:
        return json.dumps(value, ensure_ascii=False)

    text_with_placeholders = _replace_segments_with_placeholders(text=value, segments=segments)
    child_indent = " " * (indent + 2)
    current_indent = " " * indent
    lines = [f"{INLINE_STRING_COMMENT} {{"]
    lines.append(f"{child_indent}{json.dumps('_text')}: {json.dumps(text_with_placeholders, ensure_ascii=False)},")

    used_labels = {"_text"}
    for index, segment in enumerate(segments):
        label = _unique_inline_label(label=segment.label, used_labels=used_labels)
        used_labels.add(label)
        converted_value, nested_paths = _convert_json_strings_with_annotations(
            value=segment.value,
            path=(*path, label),
            max_decode_depth=DEFAULT_MAX_DECODE_DEPTH,
        )
        child_text = _format_jsonc_value(
            value=converted_value,
            decoded_string_paths=nested_paths,
            path=(*path, label),
            indent=indent + 2,
        )
        comma = "," if index < len(segments) - 1 else ""
        lines.append(
            f"{child_indent}{json.dumps(label, ensure_ascii=False)}: {INLINE_VALUE_COMMENT} {child_text}{comma}"
        )
    lines.append(f"{current_indent}}}")
    return "\n".join(lines)


def _extract_inline_json_segments(*, text: str) -> list[InlineJsonSegment]:
    decoder = json.JSONDecoder()
    segments: list[InlineJsonSegment] = []
    index = 0
    previous_end = 0

    while index < len(text):
        if text[index] not in "{[":
            index += 1
            continue

        try:
            decoded, end = decoder.raw_decode(text, index)
        except json.JSONDecodeError:
            index += 1
            continue

        if not isinstance(decoded, dict | list):
            index += 1
            continue

        segments.append(
            InlineJsonSegment(
                label=_inline_json_label(text=text[previous_end:index], fallback_index=len(segments) + 1),
                start=index,
                end=end,
                value=decoded,
            )
        )
        previous_end = end
        index = end

    return segments


def _inline_json_label(*, text: str, fallback_index: int) -> str:
    prefix = text.rstrip()
    if ":" not in prefix:
        return f"inline_json_{fallback_index}"

    label = prefix.rsplit(":", maxsplit=1)[0].rsplit(",", maxsplit=1)[-1].splitlines()[-1].strip()
    if not label or len(label) > 40:
        return f"inline_json_{fallback_index}"
    return label


def _replace_segments_with_placeholders(*, text: str, segments: list[InlineJsonSegment]) -> str:
    parts: list[str] = []
    cursor = 0
    for segment in segments:
        parts.append(text[cursor : segment.start])
        parts.append("…")
        cursor = segment.end
    parts.append(text[cursor:])
    return "".join(parts)


def _unique_inline_label(*, label: str, used_labels: set[str]) -> str:
    if label not in used_labels:
        return label

    suffix = 2
    while f"{label}_{suffix}" in used_labels:
        suffix += 1
    return f"{label}_{suffix}"


def _format_path(*, path: JSONPath) -> str:
    if not path:
        return "$"
    return "$" + "".join(f"[{part}]" if isinstance(part, int) else f".{part}" for part in path)


def _highlight_jsonc_syntax(*, text: str) -> str:
    parts: list[str] = []
    index = 0
    while index < len(text):
        if comment := _comment_at(text=text, index=index):
            parts.append(_ansi(text=comment, color=JSONX_COMMENT_HIGHLIGHT))
            index += len(comment)
            continue

        char = text[index]
        if char == '"':
            end = _json_string_end(text=text, index=index)
            token = text[index:end]
            color = JSON_KEY_COLOR if _is_object_key(text=text, index=end) else JSON_STRING_COLOR
            parts.append(_ansi(text=token, color=color))
            index = end
            continue

        if (char == "-" or char.isdigit()) and (match := JSON_NUMBER_PATTERN.match(text, index)):
            parts.append(_ansi(text=match.group(0), color=JSON_LITERAL_COLOR))
            index = match.end()
            continue

        if literal := _literal_at(text=text, index=index):
            parts.append(_ansi(text=literal, color=JSON_LITERAL_COLOR))
            index += len(literal)
            continue

        if char in "{}[]:,":
            parts.append(_ansi(text=char, color=JSON_PUNCT_COLOR))
        else:
            parts.append(char)
        index += 1
    return "".join(parts)


def _comment_at(*, text: str, index: int) -> str | None:
    for comment in (DECODED_STRING_COMMENT, INLINE_STRING_COMMENT, INLINE_VALUE_COMMENT):
        if text.startswith(comment, index):
            return comment
    return None


def _json_string_end(*, text: str, index: int) -> int:
    cursor = index + 1
    escaped = False
    while cursor < len(text):
        char = text[cursor]
        if escaped:
            escaped = False
        elif char == "\\":
            escaped = True
        elif char == '"':
            return cursor + 1
        cursor += 1
    return len(text)


def _is_object_key(*, text: str, index: int) -> bool:
    while index < len(text) and text[index].isspace():
        index += 1
    return index < len(text) and text[index] == ":"


def _literal_at(*, text: str, index: int) -> str | None:
    for literal in ("true", "false", "null"):
        end = index + len(literal)
        if text.startswith(literal, index) and (end == len(text) or not text[end].isalpha()):
            return literal
    return None


def _ansi(*, text: str, color: str) -> str:
    return f"{color}{text}{ANSI_RESET}"


def _looks_like_embedded_json(*, text: str) -> bool:
    if len(text) < 2:
        return False
    pairs = {"{": "}", "[": "]", '"': '"'}
    return pairs.get(text[0]) == text[-1]


def _skip_whitespace(*, text: str, index: int) -> int:
    while index < len(text) and text[index].isspace():
        index += 1
    return index


if __name__ == "__main__":
    raise SystemExit(main())