json_semval/jsonpath_utils.py

from __future__ import annotations

from typing import Any, List, Tuple


def path_to_jsonpath(path_tokens: List[str]) -> str:
    return "$" if not path_tokens else "$" + "".join(path_tokens)


def build_path_tokens(key_stack: List[str]) -> List[str]:
    # keys as [".user", "[0]", ".name"]
    return key_stack


def set_value_at_path(payload: Any, path_tokens: List[str], value: Any) -> Any:
    node = payload
    parents: List[Tuple[Any, str]] = []
    for token in path_tokens[:-1]:
        parents.append((node, token))
        if token.startswith("["):
            idx = int(token.strip("[]"))
            node = node[idx]
        else:
            key = token[1:] if token.startswith(".") else token
            node = node[key]
    last = path_tokens[-1]
    if last.startswith("["):
        idx = int(last.strip("[]"))
        node[idx] = value
    else:
        key = last[1:] if last.startswith(".") else last
        node[key] = value
    return payload


def get_value_at_path(payload: Any, path_tokens: List[str]) -> Any:
    node = payload
    for token in path_tokens:
        if token.startswith("["):
            idx = int(token.strip("[]"))
            node = node[idx]
        else:
            key = token[1:] if token.startswith(".") else token
            node = node[key]
    return node


def delete_key_at_path(payload: Any, path_tokens: List[str]) -> Any:
    node = payload
    for token in path_tokens[:-1]:
        if token.startswith("["):
            idx = int(token.strip("[]"))
            node = node[idx]
        else:
            key = token[1:] if token.startswith(".") else token
            node = node[key]
    last = path_tokens[-1]
    if last.startswith("["):
        idx = int(last.strip("[]"))
        del node[idx]
    else:
        key = last[1:] if last.startswith(".") else last
        if isinstance(node, dict) and key in node:
            del node[key]
    return payload