Added context manager to lruwrap.

2 files changed, 34 insertions, 1 deletions

diff --git a/lru.py b/lru.py
index 492f42a..a9d5dfa 100644
--- a/lru.py
+++ b/lru.py
@@ -328,8 +328,12 @@ class lruwrap(object):
                 self.store[key] = value
             self.dirty.clear()
             
-    def __del__(self):
+    def __enter__(self):
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
         self.sync()
+        return False
 
 
 class lrudecorator(object):
diff --git a/test.py b/test.py
index 9b58b1a..150fa60 100644
--- a/test.py
+++ b/test.py
@@ -150,6 +150,34 @@ def wraptest2():
             tmp2 = None
 
         assert tmp1 == tmp2
+
+
+def wraptest3():
+    import random
+
+    q = dict()
+    with lruwrap(q, 32, True) as x:
+        for i in range(256):
+            a = random.randint(0, 256)
+            b = random.randint(0, 256)
+        
+            x[a] = b
+        
+    for i in range(512):
+        a = random.randint(0, 256)
+        tmp1 = None
+        tmp2 = None
+        try:
+            tmp1 = x[a]
+        except KeyError:
+            tmp1 = None
+           
+        try:
+            tmp2 = q[a]
+        except KeyError:
+            tmp2 = None
+
+        assert tmp1 == tmp2
         
         
 @lrudecorator(14)
@@ -161,6 +189,7 @@ if __name__ == '__main__':
     
     wraptest()
     wraptest2()
+    wraptest3()
     
     for i in range(300):
         x = random.randint(0, 25)