patx/pickledb

add context manager for save() for batched operations closes #101

Commit d669ba4 · patx · 2025-02-08T20:36:21-05:00

Changeset

d669ba4bf447b884e6283f54a5cfbce8c5383fbf

Parents

1b837936036ffc0c91ccca77eaf95e6263cd7e1c

View source at this commit

Comments

No comments yet.

Log in to comment

Diff

diff --git a/pickledb.py b/pickledb.py
index ee617b6..e13a88b 100644
--- a/pickledb.py
+++ b/pickledb.py
@@ -63,6 +63,22 @@ class PickleDB:
         """
         return self.get(key)
 
+    def __enter__(self):
+        """
+        Enter the context manager.
+        Does nothing but return `self` for modifications.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exit the context manager.
+        Automatically saves changes if no exception occurred.
+        """
+        if exc_type is None:
+            self.save()
+        return False
+
     def _load(self):
         """
         Load data from the JSON file if it exists, or initialize an empty