Fix atomic patch application corner cases

CHANGELOG.md

@@ -5,6 +5,7 @@
 **Fixes**
 
 - Fixed lexing of quoted name selectors ending in an escaped backslash, like `$['a\\']['b']`. The lexer's look-behind for the closing quote treated a quote following an escaped backslash (`\\`) as an escaped quote (`\'`), so these selectors failed to tokenize. This also broke the RFC 9535 normalized path round-trip: normalized paths produced for object keys containing backslashes could not be parsed back. See [#132](https://github.com/jg-rp/python-jsonpath/pull/132).
+- Fixed atomic JSONPatch application when a patch op replaces the document root.
 
 **Features**
 

jsonpath/patch.py

@@ -694,16 +694,19 @@ def atomic(
             JSONPatchTestFailure: When a _test_ operation does not pass.
                 `JSONPatchTestFailure` is a subclass of `JSONPatchError`.
         """
-        data_ = copy.deepcopy(data)
-        self.apply(data_)  # This could raise a JSONPatchError.
-        data.clear()
+        data_ = self.apply(copy.deepcopy(data))  # This could raise a JSONPatchError.
 
-        if isinstance(data, dict):
-            data.update(data_)
-        else:
-            data.extend(data_)
+        if isinstance(data, dict) and isinstance(data_, dict):
+            data.clear()
+            data.update(data_)  # type: ignore
+            return data
 
-        return data
+        if isinstance(data, list) and isinstance(data_, list):
+            data.clear()
+            data.extend(data_)  # type: ignore
+            return data
+
+        return data_
 
     def patched(
         self,

tests/test_json_patch.py

@@ -354,6 +354,62 @@ def test_atomic_patch_array_fail() -> None:
     assert data == data_
 
 
+def test_atomically_replace_root() -> None:
+    ops: List[Dict[str, Any]] = [
+        {"op": "replace", "path": "", "value": {"foo": "bar"}},
+    ]
+
+    data: Dict[str, Any] = {"a": {"b": {"c": 1}}}
+    patched_data = patch.atomic(ops, data)
+
+    # Note that atomic patch application clears and updates `data` when
+    # the document root is replaced with a compatible object. `patch.apply()`
+    # does not do this.
+    assert data == patched_data
+    assert data == {"foo": "bar"}
+
+
+def test_atomically_replace_root_with_incompatible_type() -> None:
+    ops: List[Dict[str, Any]] = [
+        {"op": "replace", "path": "", "value": [1, 2, 3]},
+    ]
+
+    data: Dict[str, Any] = {"a": {"b": {"c": 1}}}
+    patched_data = patch.atomic(ops, data)
+
+    assert data != patched_data
+    assert data == {"a": {"b": {"c": 1}}}
+    assert patched_data == [1, 2, 3]
+
+
+def test_atomically_replace_root_primitive() -> None:
+    ops: List[Dict[str, Any]] = [
+        {"op": "replace", "path": "", "value": "foo"},
+    ]
+
+    data: Dict[str, Any] = {"a": {"b": {"c": 1}}}
+    patched_data = patch.atomic(ops, data)
+
+    assert data != patched_data
+    assert data == {"a": {"b": {"c": 1}}}
+    assert patched_data == "foo"
+
+
+def test_apply_replace_root() -> None:
+    ops: List[Dict[str, Any]] = [
+        {"op": "replace", "path": "", "value": {"foo": "bar"}},
+    ]
+
+    data: Dict[str, Any] = {"a": {"b": {"c": 1}}}
+    patched_data = patch.apply(ops, data)
+
+    # Note that patch application can not replace `data` when the document root
+    # is replaced.
+    assert data != patched_data
+    assert data == {"a": {"b": {"c": 1}}}
+    assert patched_data == {"foo": "bar"}
+
+
 def test_patched_does_not_mutate_data() -> None:
     """Test that _patched_ modifies a deep copy of data."""
     patch_doc: List[Dict[str, Any]] = [