Merge pull request #267 from llllllllll/memoize

Memoize perf

toolz/functoolz.py

@@ -1,4 +1,4 @@
-from functools import reduce, partial
+from functools import reduce, partial, wraps
 import inspect
 import operator
 import sys
@@ -332,27 +332,28 @@ def memoize(func, cache=None, key=None):
         may_have_kwargs = True
         is_unary = False
 
+    if key is None:
+        if is_unary:
+            def key(args, kwargs):
+                return args[0]
+        elif may_have_kwargs:
+            def key(args, kwargs):
+                return (
+                    args or None,
+                    frozenset(kwargs.items()) if kwargs else None,
+                )
+        else:
+            def key(args, kwargs):
+                return args
+
     def memof(*args, **kwargs):
+        k = key(args, kwargs)
         try:
-            if key is not None:
-                k = key(args, kwargs)
-            elif is_unary:
-                k = args[0]
-            elif may_have_kwargs:
-                k = (args or None,
-                     frozenset(kwargs.items()) if kwargs else None)
-            else:
-                k = args
-
-            in_cache = k in cache
+            return cache[k]
         except TypeError:
             raise TypeError("Arguments to memoized function must be hashable")
-
-        if in_cache:
-            return cache[k]
-        else:
-            result = func(*args, **kwargs)
-            cache[k] = result
+        except KeyError:
+            cache[k] = result = func(*args, **kwargs)
             return result
 
     try:

toolz/tests/test_functoolz.py

@@ -47,7 +47,7 @@ def f(x, y):
         return x + y
     mf = memoize(f)
 
-    assert mf(2, 3) == mf(2, 3)
+    assert mf(2, 3) is mf(2, 3)
     assert fn_calls == [1]  # function was only called once
     assert mf.__doc__ == f.__doc__
     assert raises(TypeError, lambda: mf(1, {}))