Apply JSON Patch Documents

from __future__ import annotations

import copy
import json
from typing import Any


class JsonPatchError(ValueError):
    """Raised when a JSON Patch document cannot be applied."""


def _parse_pointer(pointer: str) -> list[str]:
    if not isinstance(pointer, str):
        raise TypeError("JSON Pointer must be a string")
    if pointer == "":
        return []
    if not pointer.startswith("/"):
        raise JsonPatchError("JSON Pointer must start with '/'")
    return [
        token.replace("~1", "/").replace("~0", "~")
        for token in pointer.split("/")[1:]
    ]


def _array_index(
    token: str,
    length: int,
    *,
    allow_end: bool = False,
) -> int:
    if token == "-":
        raise JsonPatchError("'-' is only valid for add operations")
    if not token.isdigit() or (len(token) > 1 and token.startswith("0")):
        raise JsonPatchError(f"invalid array index '{token}'")
    index = int(token)
    upper = length if allow_end else length - 1
    if index > upper:
        raise JsonPatchError(f"array index {index} is out of range")
    return index


def _get(document: Any, tokens: list[str]) -> Any:
    current = document
    for token in tokens:
        if isinstance(current, list):
            current = current[_array_index(token, len(current))]
        elif isinstance(current, dict):
            if token not in current:
                raise JsonPatchError(f"path segment '{token}' does not exist")
            current = current[token]
        else:
            raise JsonPatchError(
                f"cannot traverse into {type(current).__name__}"
            )
    return current


def _add(document: Any, tokens: list[str], value: Any) -> Any:
    if not tokens:
        return value
    parent = _get(document, tokens[:-1])
    token = tokens[-1]
    if isinstance(parent, list):
        if token == "-":
            parent.append(value)
        else:
            parent.insert(
                _array_index(token, len(parent), allow_end=True),
                value,
            )
    elif isinstance(parent, dict):
        parent[token] = value
    else:
        raise JsonPatchError(f"cannot add into {type(parent).__name__}")
    return document


def _remove(document: Any, tokens: list[str]) -> tuple[Any, Any]:
    if not tokens:
        raise JsonPatchError("removing the document root is not supported")
    parent = _get(document, tokens[:-1])
    token = tokens[-1]
    if isinstance(parent, list):
        value = parent.pop(_array_index(token, len(parent)))
    elif isinstance(parent, dict):
        if token not in parent:
            raise JsonPatchError(f"path segment '{token}' does not exist")
        value = parent.pop(token)
    else:
        raise JsonPatchError(f"cannot remove from {type(parent).__name__}")
    return document, value


def _replace(document: Any, tokens: list[str], value: Any) -> Any:
    if not tokens:
        return value
    parent = _get(document, tokens[:-1])
    token = tokens[-1]
    if isinstance(parent, list):
        parent[_array_index(token, len(parent))] = value
    elif isinstance(parent, dict):
        if token not in parent:
            raise JsonPatchError(f"path segment '{token}' does not exist")
        parent[token] = value
    else:
        raise JsonPatchError(f"cannot replace in {type(parent).__name__}")
    return document


def _json_equal(left: Any, right: Any) -> bool:
    if type(left) is not type(right):
        return False
    if isinstance(left, dict):
        return left.keys() == right.keys() and all(
            _json_equal(left[key], right[key]) for key in left
        )
    if isinstance(left, list):
        return len(left) == len(right) and all(
            _json_equal(a, b) for a, b in zip(left, right)
        )
    return left == right


def apply_patch(document: Any, patch: list[dict[str, Any]]) -> Any:
    if not isinstance(patch, list):
        raise TypeError("patch must be a list of operations")
    result = copy.deepcopy(document)

    for index, operation in enumerate(patch):
        if not isinstance(operation, dict):
            raise JsonPatchError(f"operation {index} must be an object")
        op = operation.get("op")
        path = operation.get("path")
        if not isinstance(op, str) or not isinstance(path, str):
            raise JsonPatchError(
                f"operation {index} requires string op and path"
            )
        tokens = _parse_pointer(path)
        if op in {"add", "replace", "test"} and "value" not in operation:
            raise JsonPatchError(f"operation {index} requires a value")

        try:
            if op == "add":
                result = _add(
                    result,
                    tokens,
                    copy.deepcopy(operation["value"]),
                )
            elif op == "remove":
                result, _ = _remove(result, tokens)
            elif op == "replace":
                result = _replace(
                    result,
                    tokens,
                    copy.deepcopy(operation["value"]),
                )
            elif op == "move":
                from_path = operation.get("from")
                if not isinstance(from_path, str):
                    raise JsonPatchError("move requires a from path")
                from_tokens = _parse_pointer(from_path)
                if not from_tokens:
                    raise JsonPatchError("cannot move the document root")
                is_child = (
                    len(from_tokens) < len(tokens)
                    and tokens[: len(from_tokens)] == from_tokens
                )
                if is_child:
                    raise JsonPatchError(
                        "cannot move a value into its own child"
                    )
                result, moved = _remove(result, from_tokens)
                result = _add(result, tokens, moved)
            elif op == "copy":
                from_path = operation.get("from")
                if not isinstance(from_path, str):
                    raise JsonPatchError("copy requires a from path")
                value = copy.deepcopy(
                    _get(result, _parse_pointer(from_path))
                )
                result = _add(result, tokens, value)
            elif op == "test":
                if not _json_equal(_get(result, tokens), operation["value"]):
                    raise JsonPatchError(f"test failed at {path}")
            else:
                raise JsonPatchError(f"unsupported operation '{op}'")
        except JsonPatchError as error:
            raise JsonPatchError(
                f"operation {index} ({op}): {error}"
            ) from error

    return result


def main() -> None:
    document = {
        "name": "demo",
        "tags": ["a", "b"],
        "meta": {"active": True, "count": 2},
    }
    patch = [
        {"op": "replace", "path": "/name", "value": "updated"},
        {"op": "add", "path": "/tags/-", "value": "c"},
        {"op": "add", "path": "/flags", "value": {}},
        {"op": "copy", "from": "/meta/count", "path": "/meta/copiedCount"},
        {"op": "move", "from": "/meta/active", "path": "/flags/active"},
        {"op": "test", "path": "/meta/count", "value": 2},
        {"op": "remove", "path": "/meta/copiedCount"},
    ]
    updated = apply_patch(document, patch)
    print(json.dumps(updated, indent=2, sort_keys=True))
    print(json.dumps(document, sort_keys=True))

    try:
        apply_patch(
            document,
            [{"op": "test", "path": "/meta/count", "value": 3}],
        )
    except JsonPatchError as error:
        print(error)

    try:
        apply_patch(
            document,
            [{"op": "move", "from": "/meta", "path": "/meta/child"}],
        )
    except JsonPatchError as error:
        print(error)


if __name__ == "__main__":
    main()