jsonpath_cli.py

#!/usr/bin/env python3
"""
JSONPath CLI - Query JSON files using JSONPath expressions
"""

import json
import sys
import argparse
from typing import Any, List
from jsonpath_ng import parse
from jsonpath_ng.ext import parser


class Colors:
    """ANSI color codes for terminal output"""
    RESET = '\033[0m'
    BOLD = '\033[1m'
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'


def colorize_json(obj: Any, indent: int = 0) -> str:
    """Colorize JSON output for better readability"""
    spaces = "  " * indent
    
    if isinstance(obj, dict):
        if not obj:
            return "{}"
        lines = ["{"]
        items = list(obj.items())
        for i, (key, value) in enumerate(items):
            comma = "," if i < len(items) - 1 else ""
            lines.append(f'{spaces}  {Colors.CYAN}"{key}"{Colors.RESET}: {colorize_json(value, indent + 1)}{comma}')
        lines.append(f'{spaces}}}')
        return "\n".join(lines)
    
    elif isinstance(obj, list):
        if not obj:
            return "[]"
        lines = ["["]
        for i, item in enumerate(obj):
            comma = "," if i < len(obj) - 1 else ""
            lines.append(f'{spaces}  {colorize_json(item, indent + 1)}{comma}')
        lines.append(f'{spaces}]')
        return "\n".join(lines)
    
    elif isinstance(obj, str):
        return f'{Colors.GREEN}"{obj}"{Colors.RESET}'
    
    elif isinstance(obj, bool):
        return f'{Colors.YELLOW}{str(obj).lower()}{Colors.RESET}'
    
    elif obj is None:
        return f'{Colors.MAGENTA}null{Colors.RESET}'
    
    elif isinstance(obj, (int, float)):
        return f'{Colors.BLUE}{obj}{Colors.RESET}'
    
    return str(obj)


def query_json(data: Any, jsonpath_expr: str) -> List[Any]:
    """Query JSON data using JSONPath expression"""
    try:
        # Try extended parser first (supports filters)
        jsonpath_expression = parser.ExtentedJsonPathParser().parse(jsonpath_expr)
    except Exception:
        # Fall back to basic parser
        jsonpath_expression = parse(jsonpath_expr)
    
    matches = jsonpath_expression.find(data)
    return [match.value for match in matches]


def main():
    parser_args = argparse.ArgumentParser(
        description='Query JSON files using JSONPath expressions',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s data.json "$.users[*].name"
  %(prog)s data.json "$.users[?(@.age > 25)]"
  cat data.json | %(prog)s - "$.store.book[0]"
  %(prog)s data.json "$.users[*].email" --raw
        """
    )
    
    parser_args.add_argument(
        'file',
        help='JSON file to query (use "-" for stdin)'
    )
    
    parser_args.add_argument(
        'jsonpath',
        help='JSONPath expression (e.g., "$.users[*].name")'
    )
    
    parser_args.add_argument(
        '--compact', '-c',
        action='store_true',
        help='Output compact JSON (no pretty-printing)'
    )
    
    parser_args.add_argument(
        '--raw', '-r',
        action='store_true',
        help='Output raw values without JSON formatting'
    )
    
    parser_args.add_argument(
        '--no-color',
        action='store_true',
        help='Disable colored output'
    )
    
    args = parser_args.parse_args()
    
    # Disable colors if requested or not a TTY
    if args.no_color or not sys.stdout.isatty():
        for attr in dir(Colors):
            if not attr.startswith('_'):
                setattr(Colors, attr, '')
    
    try:
        # Read JSON data
        if args.file == '-':
            data = json.load(sys.stdin)
        else:
            with open(args.file, 'r') as f:
                data = json.load(f)
        
        # Query the data
        results = query_json(data, args.jsonpath)
        
        if not results:
            print(f"{Colors.YELLOW}No matches found{Colors.RESET}", file=sys.stderr)
            sys.exit(1)
        
        # Output results
        if args.raw:
            # Raw output - one value per line
            for result in results:
                if isinstance(result, (dict, list)):
                    print(json.dumps(result, ensure_ascii=False))
                else:
                    print(result)
        elif args.compact:
            # Compact JSON output
            if len(results) == 1:
                print(json.dumps(results[0], ensure_ascii=False))
            else:
                print(json.dumps(results, ensure_ascii=False))
        else:
            # Pretty-printed colored output
            if len(results) == 1:
                print(colorize_json(results[0]))
            else:
                print(colorize_json(results))
    
    except FileNotFoundError:
        print(f"{Colors.RED}Error: File '{args.file}' not found{Colors.RESET}", file=sys.stderr)
        sys.exit(1)
    
    except json.JSONDecodeError as e:
        print(f"{Colors.RED}Error: Invalid JSON - {e}{Colors.RESET}", file=sys.stderr)
        sys.exit(1)
    
    except Exception as e:
        print(f"{Colors.RED}Error: {e}{Colors.RESET}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()