func: Add lru_cache(), a coroutine version of functools.lru_cache

Author: achimnol

aiotools/func.py

@@ -1,7 +1,9 @@
+import collections
 import functools
 
 __all__ = [
     'apartial',
+    'lru_cache',
 ]
 
 
@@ -14,3 +16,64 @@ def apartial(coro, *args, **kwargs):
     async def wrapped(*cargs, **ckwargs):
         return await coro(*args, *cargs, **kwargs, **ckwargs)
     return wrapped
+
+
+def lru_cache(maxsize=128, typed=False):
+    '''
+    A simple LRU cache just like functools.lru_cache, but it works for
+    coroutines.  This is not as heavily optimized as functools.lru_cache
+    which uses an internal C implementation, as it targets async operations
+    that take a long time.
+
+    It follows the same API that the standard functools provides.  The wrapped
+    function has ``cache_clear()`` method to flush the cache manually, but
+    leaves ``cache_info()`` for statistics unimplemented.
+
+    Note that calling the coroutine multiple times with the same arguments
+    before the first call returns may incur duplicate exectuions.
+
+    This function is not thread-safe.
+    '''
+
+    if maxsize is not None and not isinstance(maxsize, int):
+        raise TypeError('Expected maxsize to be an integer or None')
+
+    def wrapper(coro):
+
+        sentinel = object()  # unique object to distinguish None as result
+        cache = collections.OrderedDict()
+        cache_get = cache.get
+        cache_set = cache.__setitem__
+        cache_len = cache.__len__
+        cache_move = cache.move_to_end
+        make_key = functools._make_key
+
+        # We don't use explicit locks like the standard functools,
+        # because this lru_cache is intended for use in asyncio coroutines.
+        # The only context interleaving happens when calling the user-defined
+        # coroutine, so there is no need to add extra synchronization guards.
+
+        @functools.wraps(coro)
+        async def wrapped(*args, **kwargs):
+            k = make_key(args, kwargs, typed)
+            result = cache_get(k, sentinel)
+            if result is not sentinel:
+                return result
+            result = await coro(*args, **kwargs)
+            if maxsize is not None and cache_len() >= maxsize:
+                cache.popitem(last=False)
+            cache_set(k, result)
+            cache_move(k, last=True)
+            return result
+
+        def cache_clear():
+            cache.clear()
+
+        def cache_info():
+            raise NotImplementedError
+
+        wrapped.cache_clear = cache_clear
+        wrapped.cache_info = cache_info
+        return wrapped
+
+    return wrapper

tests/test_func.py

@@ -1,6 +1,8 @@
+import asyncio
+
 import pytest
 
-from aiotools.func import apartial
+from aiotools.func import apartial, lru_cache
 
 
 async def do(a, b, *, c=1, d=2):
@@ -43,3 +45,43 @@ async def test_apartial_wraps():
     assert do2.__doc__ == do.__doc__
     assert do.__name__ == 'do'
     assert do2.__name__ == 'do'
+
+
+@pytest.mark.asyncio
+async def test_lru_cache():
+
+    calc_count = 0
+
+    @lru_cache(maxsize=2)
+    async def calc(n):
+        '''testing'''
+        nonlocal calc_count
+        await asyncio.sleep(0)
+        calc_count += 1
+        return n * n
+
+    assert calc.__name__ == 'calc'
+    assert calc.__doc__ == 'testing'
+
+    assert (await calc(1)) == 1
+    assert calc_count == 1
+    assert (await calc(2)) == 4
+    assert calc_count == 2
+    assert (await calc(1)) == 1
+    assert calc_count == 2
+    assert (await calc(3)) == 9
+    assert calc_count == 3
+    assert (await calc(1)) == 1  # evicted and re-executed
+    assert calc_count == 4
+    assert (await calc(1)) == 1  # cached again
+    assert calc_count == 4
+
+    with pytest.raises(NotImplementedError):
+        calc.cache_info()
+
+    calc.cache_clear()
+
+    assert (await calc(1)) == 1
+    assert calc_count == 5
+    assert (await calc(3)) == 9
+    assert calc_count == 6