|
| 1 | +import json |
| 2 | +from ast import literal_eval |
| 3 | +from importlib.util import find_spec |
| 4 | +from typing import Any, Dict |
| 5 | + |
| 6 | +if find_spec("colorama"): |
| 7 | + import colorama |
| 8 | + |
| 9 | + colorama.init() |
| 10 | + |
| 11 | + |
| 12 | +# ANSI color codes |
| 13 | +RED = '\033[31m' |
| 14 | +GREEN = '\033[32m' |
| 15 | +RESET = '\033[0m' |
| 16 | + |
| 17 | +class ColoredView: |
| 18 | + """A view that shows JSON with color-coded differences.""" |
| 19 | + |
| 20 | + def __init__(self, t2, tree_results, verbose_level=1): |
| 21 | + self.t2 = t2 |
| 22 | + self.tree = tree_results |
| 23 | + self.verbose_level = verbose_level |
| 24 | + self.diff_paths = self._collect_diff_paths() |
| 25 | + |
| 26 | + def _collect_diff_paths(self) -> Dict[str, str]: |
| 27 | + """Collect all paths that have differences and their types.""" |
| 28 | + diff_paths = {} |
| 29 | + for diff_type, items in self.tree.items(): |
| 30 | + try: |
| 31 | + iterator = iter(items) |
| 32 | + except TypeError: |
| 33 | + continue |
| 34 | + for item in items: |
| 35 | + if type(item).__name__ == "DiffLevel": |
| 36 | + path = item.path() |
| 37 | + if diff_type in ('values_changed', 'type_changes'): |
| 38 | + diff_paths[path] = ('changed', item.t1, item.t2) |
| 39 | + elif diff_type in ('dictionary_item_added', 'iterable_item_added', 'set_item_added'): |
| 40 | + diff_paths[path] = ('added', None, item.t2) |
| 41 | + elif diff_type in ('dictionary_item_removed', 'iterable_item_removed', 'set_item_removed'): |
| 42 | + diff_paths[path] = ('removed', item.t1, None) |
| 43 | + return diff_paths |
| 44 | + |
| 45 | + def _format_value(self, value: Any) -> str: |
| 46 | + """Format a value for display.""" |
| 47 | + if isinstance(value, bool): |
| 48 | + return 'true' if value else 'false' |
| 49 | + elif isinstance(value, str): |
| 50 | + return f'"{value}"' |
| 51 | + elif isinstance(value, (dict, list, tuple)): |
| 52 | + return json.dumps(value) |
| 53 | + else: |
| 54 | + return str(value) |
| 55 | + |
| 56 | + def _get_path_removed(self, path: str) -> dict: |
| 57 | + """Get all removed items for a given path.""" |
| 58 | + removed = {} |
| 59 | + for key, value in self.diff_paths.items(): |
| 60 | + if value[0] == 'removed' and key.startswith(path + "["): |
| 61 | + key_suffix = key[len(path):] |
| 62 | + if key_suffix.count("[") == 1 and key_suffix.endswith("]"): |
| 63 | + removed[literal_eval(key_suffix[1:-1])] = value[1] |
| 64 | + return removed |
| 65 | + |
| 66 | + def _colorize_json(self, obj: Any, path: str = 'root', indent: int = 0) -> str: |
| 67 | + """Recursively colorize JSON based on differences, with pretty-printing.""" |
| 68 | + INDENT = ' ' |
| 69 | + current_indent = INDENT * indent |
| 70 | + next_indent = INDENT * (indent + 1) |
| 71 | + if path in self.diff_paths and path not in self._colorize_skip_paths: |
| 72 | + diff_type, old, new = self.diff_paths[path] |
| 73 | + if diff_type == 'changed': |
| 74 | + return f"{RED}{self._format_value(old)}{RESET} -> {GREEN}{self._format_value(new)}{RESET}" |
| 75 | + elif diff_type == 'added': |
| 76 | + return f"{GREEN}{self._format_value(new)}{RESET}" |
| 77 | + elif diff_type == 'removed': |
| 78 | + return f"{RED}{self._format_value(old)}{RESET}" |
| 79 | + |
| 80 | + if isinstance(obj, dict): |
| 81 | + if not obj: |
| 82 | + return '{}' |
| 83 | + items = [] |
| 84 | + for key, value in obj.items(): |
| 85 | + new_path = f"{path}['{key}']" if isinstance(key, str) else f"{path}[{key}]" |
| 86 | + if new_path in self.diff_paths and self.diff_paths[new_path][0] == 'added': |
| 87 | + # Colorize both key and value for added fields |
| 88 | + items.append(f'{next_indent}{GREEN}"{key}": {self._colorize_json(value, new_path, indent + 1)}{RESET}') |
| 89 | + else: |
| 90 | + items.append(f'{next_indent}"{key}": {self._colorize_json(value, new_path, indent + 1)}') |
| 91 | + for key, value in self._get_path_removed(path).items(): |
| 92 | + new_path = f"{path}['{key}']" if isinstance(key, str) else f"{path}[{key}]" |
| 93 | + items.append(f'{next_indent}{RED}"{key}": {self._colorize_json(value, new_path, indent + 1)}{RESET}') |
| 94 | + return '{\n' + ',\n'.join(items) + f'\n{current_indent}' + '}' |
| 95 | + |
| 96 | + elif isinstance(obj, (list, tuple)): |
| 97 | + if not obj: |
| 98 | + return '[]' |
| 99 | + removed_map = self._get_path_removed(path) |
| 100 | + for index in removed_map: |
| 101 | + self._colorize_skip_paths.add(f"{path}[{index}]") |
| 102 | + items = [] |
| 103 | + index = 0 |
| 104 | + for value in obj: |
| 105 | + new_path = f"{path}[{index}]" |
| 106 | + while index == next(iter(removed_map), None): |
| 107 | + items.append(f'{next_indent}{RED}{self._format_value(removed_map.pop(index))}{RESET}') |
| 108 | + index += 1 |
| 109 | + items.append(f'{next_indent}{self._colorize_json(value, new_path, indent + 1)}') |
| 110 | + index += 1 |
| 111 | + for value in removed_map.values(): |
| 112 | + items.append(f'{next_indent}{RED}{self._format_value(value)}{RESET}') |
| 113 | + return '[\n' + ',\n'.join(items) + f'\n{current_indent}' + ']' |
| 114 | + else: |
| 115 | + return self._format_value(obj) |
| 116 | + |
| 117 | + def __str__(self) -> str: |
| 118 | + """Return the colorized, pretty-printed JSON string.""" |
| 119 | + self._colorize_skip_paths = set() |
| 120 | + return self._colorize_json(self.t2, indent=0) |
| 121 | + |
| 122 | + def __iter__(self): |
| 123 | + """Make the view iterable by yielding the tree results.""" |
| 124 | + yield from self.tree.items() |
0 commit comments