🐛 Make Library.remove atomic and tolerant of changed block keys

bibtexparser/library.py

@@ -76,34 +76,84 @@ def _find_duplicate_keys(self, blocks: list[Block]) -> list[str]:
                 seen_string_keys.add(block.key)
         return duplicate_keys
 
-    def _block_index(self, block: Block) -> int:
+    def _block_index(self, block: Block, skip: set[int] | None = None) -> int:
         """Index of a block in the library, preferring identity over equality.
 
         :param block: Block to look up.
+        :param skip: Indices to ignore, e.g. because they are already claimed
+            by another block of the same (multi-block) operation.
         :raises ValueError: If block is not in library."""
+        skip = set() if skip is None else skip
         for i, b in enumerate(self._blocks):
-            if b is block:
+            if i not in skip and b is block:
                 return i
-        # No identity match; fall back to equality (raises ValueError if not found).
-        return self._blocks.index(block)
+        # No identity match; fall back to equality.
+        for i, b in enumerate(self._blocks):
+            if i not in skip and b == block:
+                return i
+        raise ValueError("Block not in library.")
+
+    def _key_registration(
+        self, block: Block, claimed: list[tuple[dict, str]]
+    ) -> tuple[dict, str] | None:
+        """The by-key dict and key under which a block is registered, if any.
+
+        A block's key is mutable and may have been changed after the block was
+        added to the library. The registration is thus primarily located by
+        identity, and only then by the block's current key.
+
+        :param block: Block to look up.
+        :param claimed: Registrations already claimed by other blocks of the
+            same operation; these are not returned a second time.
+        :returns: Tuple of by-key dict and key, or None if the block is not
+            registered in any by-key dict (e.g. comments or failed blocks)."""
+        if isinstance(block, Entry):
+            by_key = self._entries_by_key
+        elif isinstance(block, String):
+            by_key = self._strings_by_key
+        else:
+            return None
+
+        def is_claimed(key: str) -> bool:
+            """Whether this key was already claimed by an earlier block in this call."""
+            return any(d is by_key and k == key for d, k in claimed)
+
+        for key, registered in by_key.items():
+            if registered is block and not is_claimed(key):
+                return by_key, key
+        if block.key in by_key and not is_claimed(block.key):
+            return by_key, block.key
+        return None
 
     def remove(self, blocks: list[Block] | Block):
         """Remove blocks from library.
 
         If equal duplicate blocks exist in the library, the exact (identical)
         instance is removed, if present; otherwise the first equal block.
 
+        Blocks whose key was changed after they were added are removed
+        correctly, i.e., they are also de-registered from the by-key
+        dictionaries backing `entries_dict` and `strings_dict`.
+
         :param blocks: Block or list of blocks to remove.
-        :raises ValueError: If block is not in library."""
+        :raises ValueError: If a block is not in library. In this case, no
+            blocks are removed from the library."""
         if isinstance(blocks, Block):
             blocks = [blocks]
 
+        # Resolve all lookups before mutating, so a failure changes nothing.
+        indices: set[int] = set()
+        registrations: list[tuple[dict, str]] = []
         for block in blocks:
-            del self._blocks[self._block_index(block)]
-            if isinstance(block, Entry):
-                del self._entries_by_key[block.key]
-            elif isinstance(block, String):
-                del self._strings_by_key[block.key]
+            indices.add(self._block_index(block, skip=indices))
+            registration = self._key_registration(block, claimed=registrations)
+            if registration is not None:
+                registrations.append(registration)
+
+        for index in sorted(indices, reverse=True):
+            del self._blocks[index]
+        for by_key, key in registrations:
+            del by_key[key]
 
     def replace(self, old_block: Block, new_block: Block, fail_on_duplicate_key: bool = True):
         """Replace a block with another block, at the same position.

tests/test_library.py

@@ -4,6 +4,7 @@
 from bibtexparser.model import Entry
 from bibtexparser.model import Field
 from bibtexparser.model import ImplicitComment
+from bibtexparser.model import String
 
 
 def get_dummy_entry():
@@ -17,6 +18,12 @@ def get_dummy_entry():
     )
 
 
+def get_dummy_entry_with_key(key: str):
+    entry = get_dummy_entry()
+    entry.key = key
+    return entry
+
+
 def test_replace_with_duplicates():
     """Test that replace() works when there are duplicate values. See issue 404."""
     library = Library()
@@ -125,3 +132,89 @@ def test_replace_prefers_identical_instance_over_equal_block():
     assert len(library.blocks) == 2
     assert library.blocks[0] is first_comment
     assert library.blocks[1] is new_comment
+
+
+def test_remove_entry_with_changed_key():
+    """Removing an entry whose key changed after adding must not fail. See issue 565."""
+    library = Library()
+    entry = get_dummy_entry()
+    other_entry = get_dummy_entry()
+    other_entry.key = "otherKey"
+    library.add([entry, other_entry])
+
+    entry.key = "changedKey"
+    library.remove(entry)
+
+    assert library.blocks == [other_entry]
+    assert library.entries == [other_entry]
+    assert library.entries_dict == {"otherKey": other_entry}
+    library.add(get_dummy_entry())
+    assert len(library.blocks) == 2
+    assert len(library.failed_blocks) == 0
+
+
+def test_remove_string_with_changed_key():
+    """Removing a string whose key changed after adding must not fail."""
+    library = Library()
+    string = String(key="someString", value='"some value"')
+    other_string = String(key="otherString", value='"other value"')
+    library.add([string, other_string])
+
+    string.key = "changedKey"
+    library.remove(string)
+
+    assert library.blocks == [other_string]
+    assert library.strings == [other_string]
+    assert library.strings_dict == {"otherString": other_string}
+
+
+def test_remove_block_not_in_library_raises_value_error():
+    library = Library()
+    entry = get_dummy_entry()
+    library.add(entry)
+
+    with pytest.raises(ValueError):
+        library.remove(get_dummy_entry_with_key("notInLibrary"))
+
+    assert library.blocks == [entry]
+    assert library.entries_dict == {"duplicateKey": entry}
+
+
+def test_remove_list_with_missing_block_is_atomic():
+    """If one block of a removed list is missing, no block is removed."""
+    library = Library()
+    entry = get_dummy_entry()
+    library.add(entry)
+
+    with pytest.raises(ValueError):
+        library.remove([entry, get_dummy_entry_with_key("notInLibrary")])
+
+    assert library.blocks == [entry]
+    assert library.entries == [entry]
+    assert library.entries_dict == {"duplicateKey": entry}
+
+
+def test_remove_list_of_blocks():
+    library = Library()
+    entries = [get_dummy_entry_with_key(f"key{i}") for i in range(3)]
+    library.add(entries)
+
+    library.remove([entries[0], entries[2]])
+
+    assert library.blocks == [entries[1]]
+    assert library.entries_dict == {"key1": entries[1]}
+
+
+def test_replace_entry_with_changed_key():
+    """Replacing an entry whose key changed after adding must not fail. See issue 565."""
+    library = Library()
+    entry = get_dummy_entry()
+    library.add(entry)
+    entry.key = "changedKey"
+
+    replacement = get_dummy_entry_with_key("replacementKey")
+    library.replace(entry, replacement)
+
+    assert library.blocks == [replacement]
+    assert library.entries == [replacement]
+    assert library.entries_dict == {"replacementKey": replacement}