Fix #163: Support user_function with cachetools.func decorators.

3 files changed, 42 insertions, 14 deletions

diff --git a/cachetools/func.py b/cachetools/func.py
index 581877b..0124a49 100644
--- a/cachetools/func.py
+++ b/cachetools/func.py
@@ -93,6 +93,8 @@ def lfu_cache(maxsize=128, typed=False):
     """
     if maxsize is None:
         return _cache(_UnboundCache(), typed)
+    elif callable(maxsize):
+        return _cache(LFUCache(128), typed)(maxsize)
     else:
         return _cache(LFUCache(maxsize), typed)
 
@@ -105,6 +107,8 @@ def lru_cache(maxsize=128, typed=False):
     """
     if maxsize is None:
         return _cache(_UnboundCache(), typed)
+    elif callable(maxsize):
+        return _cache(LRUCache(128), typed)(maxsize)
     else:
         return _cache(LRUCache(maxsize), typed)
 
@@ -117,6 +121,8 @@ def rr_cache(maxsize=128, choice=random.choice, typed=False):
     """
     if maxsize is None:
         return _cache(_UnboundCache(), typed)
+    elif callable(maxsize):
+        return _cache(RRCache(128, choice), typed)(maxsize)
     else:
         return _cache(RRCache(maxsize, choice), typed)
 
@@ -128,5 +134,7 @@ def ttl_cache(maxsize=128, ttl=600, timer=time.monotonic, typed=False):
     """
     if maxsize is None:
         return _cache(_UnboundTTLCache(ttl, timer), typed)
+    elif callable(maxsize):
+        return _cache(TTLCache(128, ttl, timer), typed)(maxsize)
     else:
         return _cache(TTLCache(maxsize, ttl, timer), typed)
diff --git a/docs/index.rst b/docs/index.rst
index 257177e..3661426 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -426,31 +426,44 @@ arguments of different types will be cached separately.  For example,
 ``f(3)`` and ``f(3.0)`` will be treated as distinct calls with
 distinct results.
 
+If a `user_function` is specified instead, it must be a callable.
+This allows the decorator to be applied directly to a user function,
+leaving the `maxsize` at its default value of 128::
+
+  @cachetools.func.lru_cache
+  def count_vowels(sentence):
+      sentence = sentence.casefold()
+      return sum(sentence.count(vowel) for vowel in 'aeiou')
+
 The wrapped function is instrumented with :func:`cache_info` and
 :func:`cache_clear` functions to provide information about cache
 performance and clear the cache.  Please see the
 :func:`functools.lru_cache` documentation for details.  Also note that
 all the decorators in this module are thread-safe by default.
 
-.. decorator:: lfu_cache(maxsize=128, typed=False)
+.. decorator:: lfu_cache(user_function)
+               lfu_cache(maxsize=128, typed=False)
 
    Decorator that wraps a function with a memoizing callable that
    saves up to `maxsize` results based on a Least Frequently Used
    (LFU) algorithm.
 
-.. decorator:: lru_cache(maxsize=128, typed=False)
+.. decorator:: lru_cache(user_function)
+               lru_cache(maxsize=128, typed=False)
 
    Decorator that wraps a function with a memoizing callable that
    saves up to `maxsize` results based on a Least Recently Used (LRU)
    algorithm.
 
-.. decorator:: rr_cache(maxsize=128, choice=random.choice, typed=False)
+.. decorator:: rr_cache(user_function)
+               rr_cache(maxsize=128, choice=random.choice, typed=False)
 
    Decorator that wraps a function with a memoizing callable that
    saves up to `maxsize` results based on a Random Replacement (RR)
    algorithm.
 
-.. decorator:: ttl_cache(maxsize=128, ttl=600, timer=time.monotonic, typed=False)
+.. decorator:: ttl_cache(user_function)
+               ttl_cache(maxsize=128, ttl=600, timer=time.monotonic, typed=False)
 
    Decorator to wrap a function with a memoizing callable that saves
    up to `maxsize` results based on a Least Recently Used (LRU)
diff --git a/tests/test_func.py b/tests/test_func.py
index 473d9c9..31e96e6 100644
--- a/tests/test_func.py
+++ b/tests/test_func.py
@@ -5,8 +5,8 @@ import cachetools.func
 
 class DecoratorTestMixin(object):
 
-    def decorator(self, maxsize, typed=False, lock=None):
-        raise NotImplementedError
+    def decorator(self, maxsize, **kwargs):
+        return self.DECORATOR(maxsize, **kwargs)
 
     def test_decorator(self):
         cached = self.decorator(maxsize=2)(lambda n: n)
@@ -65,26 +65,33 @@ class DecoratorTestMixin(object):
         self.assertEqual(cached(1.0), 1.0)
         self.assertEqual(cached.cache_info(), (2, 2, 2, 2))
 
+    def test_decorator_user_function(self):
+        cached = self.decorator(lambda n: n)
+
+        self.assertEqual(cached.cache_info(), (0, 0, 128, 0))
+        self.assertEqual(cached(1), 1)
+        self.assertEqual(cached.cache_info(), (0, 1, 128, 1))
+        self.assertEqual(cached(1), 1)
+        self.assertEqual(cached.cache_info(), (1, 1, 128, 1))
+        self.assertEqual(cached(1.0), 1.0)
+        self.assertEqual(cached.cache_info(), (2, 1, 128, 1))
+
 
 class LFUDecoratorTest(unittest.TestCase, DecoratorTestMixin):
 
-    def decorator(self, maxsize, **kwargs):
-        return cachetools.func.lfu_cache(maxsize, **kwargs)
+    DECORATOR = staticmethod(cachetools.func.lfu_cache)
 
 
 class LRUDecoratorTest(unittest.TestCase, DecoratorTestMixin):
 
-    def decorator(self, maxsize, **kwargs):
-        return cachetools.func.lru_cache(maxsize, **kwargs)
+    DECORATOR = staticmethod(cachetools.func.lru_cache)
 
 
 class RRDecoratorTest(unittest.TestCase, DecoratorTestMixin):
 
-    def decorator(self, maxsize, **kwargs):
-        return cachetools.func.rr_cache(maxsize, **kwargs)
+    DECORATOR = staticmethod(cachetools.func.rr_cache)
 
 
 class TTLDecoratorTest(unittest.TestCase, DecoratorTestMixin):
 
-    def decorator(self, maxsize, **kwargs):
-        return cachetools.func.ttl_cache(maxsize, **kwargs)
+    DECORATOR = staticmethod(cachetools.func.ttl_cache)