Implement Atomic JSON Patch Application in Python

import copy

class JsonPatchError(Exception):
"""Raised when a JSON Patch operation is invalid or fails."""
pass

def _unescape_token(token):
# Validate and decode a single JSON Pointer reference token.
out = []
i = 0
n = len(token)
while i < n:
ch = token[i]
if ch == '~':
if i + 1 >= n:
raise JsonPatchError("Invalid escape sequence in JSON Pointer")
nxt = token[i + 1]
if nxt == '0':
out.append('~')
elif nxt == '1':
out.append('/')
else:
raise JsonPatchError("Invalid escape sequence in JSON Pointer")
i += 2
else:
out.append(ch)
i += 1
return ''.join(out)

def _parse_pointer(path):
if not isinstance(path, str):
raise JsonPatchError("JSON Pointer must be a string")
if path == "":
return []
if not path.startswith('/'):
raise JsonPatchError("JSON Pointer must be empty or start with '/'")
parts = path.split('/')[1:]
return [_unescape_token(p) for p in parts]

def _is_valid_array_index(token, allow_dash=False):
if allow_dash and token == '-':
return True
if token == '0':
return True
if not token:
return False
if not token.isdigit():
return False
if token[0] == '0':
return False
return True

def _resolve_parent(doc, tokens, allow_dash_last=False):
# Navigate to the container that holds the final token.
# Returns (parent, final_token).
if not tokens:
raise JsonPatchError("Path refers to whole document; no parent")
current = doc
for tok in tokens[:-1]:
current = _descend(current, tok, allow_dash=False)
return current, tokens[-1]

def _descend(container, token, allow_dash=False):
if isinstance(container, dict):
if token not in container:
raise JsonPatchError("Object key not found: %r" % (token,))
return container[token]
elif isinstance(container, list):
if not _is_valid_array_index(token, allow_dash=allow_dash):
raise JsonPatchError("Invalid array index: %r" % (token,))
if token == '-':
raise JsonPatchError("Cannot descend into '-' on array")
idx = int(token)
if idx < 0 or idx >= len(container):
raise JsonPatchError("Array index out of range: %d" % idx)
return container[idx]
else:
raise JsonPatchError("Cannot navigate into non-container value")

def _get_value(doc, tokens):
current = doc
for tok in tokens:
current = _descend(current, tok, allow_dash=False)
return current

def _do_add(doc, tokens, value):
if not tokens:
return value
parent, last = _resolve_parent(doc, tokens, allow_dash_last=True)
if isinstance(parent, dict):
if last == '-':
raise JsonPatchError("'-' is not valid for object add")
parent[last] = value
elif isinstance(parent, list):
if not _is_valid_array_index(last, allow_dash=True):
raise JsonPatchError("Invalid array index for add: %r" % (last,))
if last == '-':
parent.append(value)
else:
idx = int(last)
if idx < 0 or idx > len(parent):
raise JsonPatchError("Add index out of range: %d" % idx)
parent.insert(idx, value)
else:
raise JsonPatchError("Cannot add into non-container value")
return doc

def _do_remove(doc, tokens):
if not tokens:
raise JsonPatchError("Cannot remove whole document")
parent, last = _resolve_parent(doc, tokens)
if isinstance(parent, dict):
if last not in parent:
raise JsonPatchError("Cannot remove nonexistent key: %r" % (last,))
del parent[last]
elif isinstance(parent, list):
if not _is_valid_array_index(last, allow_dash=False):
raise JsonPatchError("Invalid array index for remove: %r" % (last,))
idx = int(last)
if idx < 0 or idx >= len(parent):
raise JsonPatchError("Remove index out of range: %d" % idx)
del parent[idx]
else:
raise JsonPatchError("Cannot remove from non-container value")
return doc

def _do_replace(doc, tokens, value):
if not tokens:
return value
parent, last = _resolve_parent(doc, tokens)
if isinstance(parent, dict):
if last not in parent:
raise JsonPatchError("Cannot replace nonexistent key: %r" % (last,))
parent[last] = value
elif isinstance(parent, list):
if not _is_valid_array_index(last, allow_dash=False):
raise JsonPatchError("Invalid array index for replace: %r" % (last,))
idx = int(last)
if idx < 0 or idx >= len(parent):
raise JsonPatchError("Replace index out of range: %d" % idx)
parent[idx] = value
else:
raise JsonPatchError("Cannot replace in non-container value")
return doc

def _deep_equal(a, b):
# Exact type-aware equality. Disallow bool/number cross matches.
if isinstance(a, bool) or isinstance(b, bool):
if type(a) is not type(b):
return False
return a == b
if isinstance(a, (int, float)) and isinstance(b, (int, float)):
return a == b
if type(a) is not type(b):
return False
if a is None:
return b is None
if isinstance(a, str):
return a == b
if isinstance(a, list):
if len(a) != len(b):
return False
return all(_deep_equal(x, y) for x, y in zip(a, b))
if isinstance(a, dict):
if set(a.keys()) != set(b.keys()):
return False
return all(_deep_equal(a[k], b[k]) for k in a)
return a == b

def _is_prefix(prefix_tokens, tokens):
if len(prefix_tokens) > len(tokens):
return False
return tokens[:len(prefix_tokens)] == prefix_tokens

_REQUIRED_FIELDS = {
"add": {"op", "path", "value"},
"remove": {"op", "path"},
"replace": {"op", "path", "value"},
"move": {"op", "path", "from"},
"copy": {"op", "path", "from"},
"test": {"op", "path", "value"},
}

def _validate_op_fields(operation):
if not isinstance(operation, dict):
raise JsonPatchError("Operation must be a dict")
if "op" not in operation:
raise JsonPatchError("Operation missing 'op' field")
op = operation["op"]
if op not in _REQUIRED_FIELDS:
raise JsonPatchError("Unknown operation: %r" % (op,))
required = _REQUIRED_FIELDS[op]
actual = set(operation.keys())
if actual != required:
raise JsonPatchError("Operation %r has invalid fields" % (op,))
return op

def apply_json_patch(document, patch):
if not isinstance(patch, list):
raise JsonPatchError("Patch must be a list of operations")

# Work on a deep copy so the original is never mutated.
working = copy.deepcopy(document)

for operation in patch:
    op = _validate_op_fields(operation)

    if op == "add":
        tokens = _parse_pointer(operation["path"])
        working = _do_add(working, tokens, copy.deepcopy(operation["value"]))

    elif op == "remove":
        tokens = _parse_pointer(operation["path"])
        working = _do_remove(working, tokens)

    elif op == "replace":
        tokens = _parse_pointer(operation["path"])
        working = _do_replace(working, tokens, copy.deepcopy(operation["value"]))

    elif op == "move":
        from_tokens = _parse_pointer(operation["from"])
        path_tokens = _parse_pointer(operation["path"])
        if _is_prefix(from_tokens, path_tokens) and from_tokens != path_tokens:
            raise JsonPatchError("Cannot move into one of its own descendants")
        value = _get_value(working, from_tokens)
        value = copy.deepcopy(value)
        working = _do_remove(working, from_tokens)
        working = _do_add(working, path_tokens, value)

    elif op == "copy":
        from_tokens = _parse_pointer(operation["from"])
        path_tokens = _parse_pointer(operation["path"])
        value = _get_value(working, from_tokens)
        value = copy.deepcopy(value)
        working = _do_add(working, path_tokens, value)

    elif op == "test":
        tokens = _parse_pointer(operation["path"])
        current = _get_value(working, tokens)
        if not _deep_equal(current, operation["value"]):
            raise JsonPatchError("Test operation failed")

return working

import copy

class JsonPatchError(Exception):
"""Custom exception for JSON Patch errors."""
pass

def _decode_token(token: str) -> str:
"""Decodes a JSON Pointer token, handling ~0 and ~1 escapes."""
if '~' in token and token.replace('~0', '').replace('~1', '').find('~') != -1:
raise JsonPatchError(f"Invalid escape sequence in token: {token}")
return token.replace('~1', '/').replace('~0', '~')

def _parse_pointer(path: str) -> list[str]:
"""Parses a JSON Pointer string into a list of decoded tokens."""
if not isinstance(path, str):
raise JsonPatchError("JSON Pointer path must be a string.")
if path == "":
return []
if not path.startswith('/'):
raise JsonPatchError(f"Invalid JSON Pointer path: {path}")
return [_decode_token(token) for token in path[1:].split('/')]

def _parse_array_index_token(token: str) -> int:
"""Parses a string token into a valid array index integer."""
if not token.isdigit():
raise JsonPatchError(f"Invalid array index: '{token}'")
if len(token) > 1 and token.startswith('0'):
raise JsonPatchError(f"Invalid array index with leading zero: '{token}'")
return int(token)

def _get_parent_and_key(doc, tokens: list[str]):
"""
Traverses the document to find the parent container and the final key/index.
Returns (parent_container, final_token).
Raises JsonPatchError if the path is invalid.
"""
if not tokens:
raise JsonPatchError("Cannot get parent of the root document.")

parent = doc
for token in tokens[:-1]:
    if isinstance(parent, dict):
        if token not in parent:
            raise JsonPatchError(f"Path segment not found: '{token}'")
        parent = parent[token]
    elif isinstance(parent, list):
        try:
            index = _parse_array_index_token(token)
            parent = parent[index]
        except (JsonPatchError, IndexError):
            raise JsonPatchError(f"Path segment not found: '{token}'")
    else:
        raise JsonPatchError(f"Cannot traverse non-container of type {type(parent).__name__}.")
        
return parent, tokens[-1]

def _get_value_at_path(doc, path_str: str):
"""
Retrieves the value at a given JSON Pointer path.
Raises JsonPatchError if the path does not exist.
"""
tokens = _parse_pointer(path_str)
if not tokens:
return doc

try:
    parent, key = _get_parent_and_key(doc, tokens)
    if isinstance(parent, dict):
        return parent[key]
    elif isinstance(parent, list):
        index = _parse_array_index_token(key)
        return parent[index]
    else:
        raise JsonPatchError("Invalid path traversal.")
except (KeyError, IndexError, JsonPatchError):
    raise JsonPatchError(f"Path does not exist: {path_str}")

def _validate_op_keys(op_dict, required_keys, op_index):
"""Validates that the operation dictionary has exactly the required keys."""
op_keys = set(op_dict.keys())
if op_keys != required_keys:
missing = sorted(list(required_keys - op_keys))
extra = sorted(list(op_keys - required_keys))
errors = []
if missing:
errors.append(f"missing keys: {missing}")
if extra:
errors.append(f"unknown keys: {extra}")
raise JsonPatchError(f"Invalid operation object at index {op_index}: {', '.join(errors)}")

def apply_json_patch(document, patch):
"""
Applies a JSON Patch-style sequence of operations to a JSON-compatible value.

The implementation does not mutate the original document. If any operation is
invalid, it raises a JsonPatchError and leaves the document unchanged.

:param document: The JSON-compatible document (dict, list, etc.).
:param patch: A list of JSON Patch operation dictionaries.
:return: A new, patched document.
:raises JsonPatchError: If any operation is invalid.
"""
if not isinstance(patch, list):
    raise JsonPatchError("Patch must be a list of operations.")

new_doc = copy.deepcopy(document)

for i, op_dict in enumerate(patch):
    if not isinstance(op_dict, dict):
        raise JsonPatchError(f"Operation at index {i} is not a dictionary.")

    op = op_dict.get('op')
    if op not in ('add', 'remove', 'replace', 'move', 'copy', 'test'):
        raise JsonPatchError(f"Unknown or missing 'op' at index {i}.")

    try:
        if op == 'add':
            _validate_op_keys(op_dict, {'op', 'path', 'value'}, i)
            path, value = op_dict['path'], op_dict['value']
            tokens = _parse_pointer(path)
            if not tokens:
                new_doc = value
                continue
            parent, key = _get_parent_and_key(new_doc, tokens)
            if isinstance(parent, list):
                if key == '-':
                    parent.append(value)
                else:
                    index = _parse_array_index_token(key)
                    if index > len(parent):
                        raise JsonPatchError(f"Index {index} out of bounds for add.")
                    parent.insert(index, value)
            elif isinstance(parent, dict):
                parent[key] = value
            else:
                raise JsonPatchError("Target for 'add' is not a list or a dictionary.")

        elif op == 'remove':
            _validate_op_keys(op_dict, {'op', 'path'}, i)
            path = op_dict['path']
            tokens = _parse_pointer(path)
            if not tokens:
                raise JsonPatchError("Cannot remove the root document.")
            parent, key = _get_parent_and_key(new_doc, tokens)
            if isinstance(parent, list):
                index = _parse_array_index_token(key)
                if index >= len(parent):
                    raise JsonPatchError("Index out of bounds for remove.")
                del parent[index]
            elif isinstance(parent, dict):
                if key not in parent:
                    raise JsonPatchError("Key not found for remove.")
                del parent[key]
            else:
                raise JsonPatchError("Target for 'remove' is not a list or a dictionary.")

        elif op == 'replace':
            _validate_op_keys(op_dict, {'op', 'path', 'value'}, i)
            path, value = op_dict['path'], op_dict['value']
            tokens = _parse_pointer(path)
            if not tokens:
                new_doc = value
                continue
            parent, key = _get_parent_and_key(new_doc, tokens)
            if isinstance(parent, list):
                index = _parse_array_index_token(key)
                if index >= len(parent):
                    raise JsonPatchError("Index out of bounds for replace.")
                parent[index] = value
            elif isinstance(parent, dict):
                if key not in parent:
                    raise JsonPatchError("Key not found for replace.")
                parent[key] = value
            else:
                raise JsonPatchError("Target for 'replace' is not a list or a dictionary.")

        elif op == 'test':
            _validate_op_keys(op_dict, {'op', 'path', 'value'}, i)
            path, expected_value = op_dict['path'], op_dict['value']
            actual_value = _get_value_at_path(new_doc, path)
            if actual_value != expected_value:
                raise JsonPatchError("Value mismatch.")

        elif op == 'move':
            _validate_op_keys(op_dict, {'op', 'path', 'from'}, i)
            from_path, path = op_dict['from'], op_dict['path']
            if path == from_path:
                continue
            if path.startswith(from_path + '/') and from_path:
                raise JsonPatchError("Cannot move an object into one of its own descendants.")
            
            from_tokens = _parse_pointer(from_path)
            if not from_tokens:
                raise JsonPatchError("Cannot move the root document.")
            from_parent, from_key = _get_parent_and_key(new_doc, from_tokens)
            if isinstance(from_parent, dict):
                if from_key not in from_parent:
                    raise JsonPatchError(f"'from' path does not exist: {from_path}")
                value_to_move = from_parent.pop(from_key)
            elif isinstance(from_parent, list):
                from_index = _parse_array_index_token(from_key)
                if from_index >= len(from_parent):
                    raise JsonPatchError(f"'from' path does not exist: {from_path}")
                value_to_move = from_parent.pop(from_index)
            else:
                raise JsonPatchError("Source for 'move' is not a list or a dictionary.")

            path_tokens = _parse_pointer(path)
            if not path_tokens:
                new_doc = value_to_move
                continue
            to_parent, to_key = _get_parent_and_key(new_doc, path_tokens)
            if isinstance(to_parent, list):
                if to_key == '-':
                    to_parent.append(value_to_move)
                else:
                    to_index = _parse_array_index_token(to_key)
                    if to_index > len(to_parent):
                        raise JsonPatchError(f"Index {to_index} out of bounds for move destination.")
                    to_parent.insert(to_index, value_to_move)
            elif isinstance(to_parent, dict):
                to_parent[to_key] = value_to_move
            else:
                raise JsonPatchError("Destination for 'move' is not a list or a dictionary.")

        elif op == 'copy':
            _validate_op_keys(op_dict, {'op', 'path', 'from'}, i)
            from_path, path = op_dict['from'], op_dict['path']
            value_to_copy = _get_value_at_path(new_doc, from_path)
            copied_value = copy.deepcopy(value_to_copy)

            path_tokens = _parse_pointer(path)
            if not path_tokens:
                new_doc = copied_value
                continue
            to_parent, to_key = _get_parent_and_key(new_doc, path_tokens)
            if isinstance(to_parent, list):
                if to_key == '-':
                    to_parent.append(copied_value)
                else:
                    to_index = _parse_array_index_token(to_key)
                    if to_index > len(to_parent):
                        raise JsonPatchError(f"Index {to_index} out of bounds for copy destination.")
                    to_parent.insert(to_index, copied_value)
            elif isinstance(to_parent, dict):
                to_parent[to_key] = copied_value
            else:
                raise JsonPatchError("Destination for 'copy' is not a list or a dictionary.")
    
    except JsonPatchError as e:
        raise JsonPatchError(f"Operation at index {i} ({op}) failed: {e}") from e
    except (KeyError, IndexError) as e:
        raise JsonPatchError(f"Operation at index {i} ({op}) failed: path does not exist.") from e

return new_doc