Fixed a few bugs and added some testing code.

2 files changed, 54 insertions, 16 deletions

diff --git a/lru.py b/lru.py
index ec257c3..6df1507 100644
--- a/lru.py
+++ b/lru.py
@@ -263,7 +263,6 @@ class lruwrap(object):
         self.store.clear()
 
     def __contains__(self, key):
-        
         if key in self.cache:
             return True
         if key in self.store:
@@ -292,17 +291,17 @@ class lruwrap(object):
 
 
 class lrudecorator(object):
-    def __init__(self, func, size):
-        self.func = func
+    def __init__(self, size):
         self.cache = lrucache(size)
         
-    def __call__(self, *args, **kwargs):
-        try:
-            value = self.cache[(args, kwargs)]
-        except KeyError:
-            pass
-        
-        value = self.func(*args, **kwargs)
-        self.cache[(args, kwargs)] = value
-        return value
-        
+    def __call__(self, func):
+        def wrapped(*args):  # XXX What about kwargs
+            try:
+                value = self.cache[args]
+            except KeyError:
+                pass
+            
+            value = func(*args)
+            self.cache[args] = value
+            return value
+        return wrapped
diff --git a/test.py b/test.py
index e0e147c..1c0a94d 100644
--- a/test.py
+++ b/test.py
@@ -1,6 +1,6 @@
 
 
-from lru import lrucache
+from lru import *
 
 def _selftest():
     
@@ -94,11 +94,50 @@ def testb():
   
     for i in range(len(vect)):
         a[vect[i]] = 0
-    
+
+
+def wraptest():
+    import random
+
+    q = dict()
+    x = lruwrap(q, 32)
+    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)
+def cube(x):
+    return x*x*x
     
 if __name__ == '__main__':
-  
     import random
+    
+    wraptest()
+    
+    for i in range(300):
+        x = random.randint(0, 25)
+        assert cube(x) == x**3
+    
+  
+
   
     a = lrucache(20)
     b = simplelrucache(20)