jq.py

import re, json

import re


def jq(json_data, path, value=None, remove=False):
    """
    Find JSON nodes matching the given path in jq-like syntax and update them with value.

    Args:
    - json_data: JSON data (dict or list)
    - path: Path to search for in jq-like syntax
    - value: The value to update matching nodes with
    - remove: Whether to remove the matching nodes

    Returns:
    - Matching node or value
    """

    def traverse_and_update(node, keys, retrieve=False):
        if not keys:
            return node

        key = keys.pop(0)
        if key == '[]':
            if isinstance(node, list):
                return [traverse_and_update(item, list(keys), retrieve) for item in node]
            else:
                return None
        elif key.startswith('[') and key.endswith(']'):
            index = int(key[1:-1])
            if isinstance(node, list) and 0 <= index < len(node):
                return traverse_and_update(node[index], keys, retrieve)
            else:
                raise KeyError
        elif isinstance(node, dict) and key in node:
            if not keys:  # Last key, update value or retrieve
                if retrieve:
                    return node[key]
                else:
                    if value is not None:
                        node[key] = value
                    if remove:
                        del node[key]
            else:
                return traverse_and_update(node[key], keys, retrieve)
        return None

    keys = re.findall(r'\w+|\[\d+\]|\[\]', path)
    if value is None and remove is False:
        return traverse_and_update(json_data, keys, retrieve=True)
    else:
        traverse_and_update(json_data, keys)
        return json_data


class FilterModule(object):
    def filters(self):
        return {
            'jq': jq,
        }


if __name__ == '__main__':
    # Example JSON data
    json_data = {
        "store": {
            "book": [
                {"title": "Book 1", "author": "Author 1"},
                {"title": "Book 2", "author": "Author 2"}
            ],
            "bicycle": {"color": "red", "price": 19.95}
        }
    }

    # Example usage
    print(jq(json_data, 'store.bicycle', remove=True))

    matching_node = jq(json_data, "store.book[].title", {"f": "new title"})
    print("Matching node:", matching_node)
    print("Updated JSON data:", json.dumps(json_data, indent=2))