Handle exceptions in batch context manager

reccmp/isledecomp/compare/db.py

@@ -165,6 +165,15 @@ def __init__(self, backref: "EntityDb") -> None:
         self._orig_to_recomp = {}
         self._recomp_to_orig = {}
 
+    def reset(self):
+        """Clear all pending changes"""
+        self._orig_insert.clear()
+        self._recomp_insert.clear()
+        self._orig.clear()
+        self._recomp.clear()
+        self._orig_to_recomp.clear()
+        self._recomp_to_orig.clear()
+
     def insert_orig(self, addr: int, **kwargs):
         self._orig_insert.setdefault(addr, {}).update(kwargs)
 
@@ -201,20 +210,16 @@ def commit(self):
         if self._orig_to_recomp:
             self.base.bulk_match(self._orig_to_recomp.items())
 
-        self._orig_insert.clear()
-        self._recomp_insert.clear()
-
-        self._orig.clear()
-        self._recomp.clear()
-
-        self._orig_to_recomp.clear()
-        self._recomp_to_orig.clear()
+        self.reset()
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-        self.commit()
+        if exc_type is not None:
+            self.reset()
+        else:
+            self.commit()
 
 
 class EntityDb:

tests/test_compare_db.py

@@ -282,3 +282,33 @@ def test_batch_match_repeat_recomp_addr(db):
 
     assert db.get_by_recomp(200).orig_addr == 101
     assert db.get_by_orig(100) is None
+
+
+def test_batch_exception_uncaught(db):
+    """When using batch context manager, an uncaught exception should clear the staged changes."""
+    try:
+        with db.batch() as batch:
+            batch.set_orig(100, name="Test")
+            batch.set_recomp(200, test=123)
+            batch.match(100, 200)
+            _ = 1 / 0
+    except ZeroDivisionError:
+        pass
+
+    assert db.get_by_orig(100) is None
+    assert db.get_by_orig(200) is None
+
+
+def test_batch_exception_caught(db):
+    """If the exception is caught, allow the batch to go through."""
+    with db.batch() as batch:
+        batch.set_orig(100, name="Test")
+        batch.set_recomp(200, test=123)
+        batch.match(100, 200)
+        try:
+            _ = 1 / 0
+        except ZeroDivisionError:
+            pass
+
+    assert db.get_by_orig(100) is not None
+    assert db.get_by_recomp(200) is not None