Fix cache_clear() when pending tasks (#506)

Author: Dreamsorcerer

async_lru/__init__.py

@@ -119,6 +119,10 @@ def cache_invalidate(self, /, *args: Hashable, **kwargs: Any) -> bool:
     def cache_clear(self) -> None:
         self.__hits = 0
         self.__misses = 0
+
+        for c in self.__cache.values():
+            if c.later_call:
+                c.later_call.cancel()
         self.__cache.clear()
         self.__tasks.clear()
 
@@ -162,7 +166,7 @@ def _cache_miss(self, key: Hashable) -> None:
     def _task_done_callback(
         self, fut: "asyncio.Future[_R]", key: Hashable, task: "asyncio.Task[_R]"
     ) -> None:
-        self.__tasks.remove(task)
+        self.__tasks.discard(task)
 
         cache_item = self.__cache.get(key)
         if self.__ttl is not None and cache_item is not None:

setup.cfg

@@ -53,7 +53,8 @@ install_requires =
 [flake8]
 exclude = .git,.env,__pycache__,.eggs
 max-line-length = 88
-ignore = N801,N802,N803,E252,W503,E133,E203
+extend-select = B950
+ignore = N801,N802,N803,E252,W503,E133,E203,E501
 
 [coverage:run]
 branch = True

tests/test_basic.py

@@ -143,9 +143,9 @@ async def coro(val: int) -> int:
         coro1._LRUCacheWrapper__cache[1].fut.result()  # type: ignore[attr-defined]
         == coro2._LRUCacheWrapper__cache[1].fut.result()  # type: ignore[attr-defined]
     )
-    assert coro1._LRUCacheWrapper__cache != coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]  # noqa: E501
-    assert coro1._LRUCacheWrapper__cache.keys() == coro2._LRUCacheWrapper__cache.keys()  # type: ignore[attr-defined]  # noqa: E501
-    assert coro1._LRUCacheWrapper__cache is not coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]  # noqa: E501
+    assert coro1._LRUCacheWrapper__cache != coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]
+    assert coro1._LRUCacheWrapper__cache.keys() == coro2._LRUCacheWrapper__cache.keys()  # type: ignore[attr-defined]
+    assert coro1._LRUCacheWrapper__cache is not coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]
 
 
 async def test_alru_cache_parameters() -> None:

tests/test_cache_clear.py

@@ -18,3 +18,38 @@ async def coro(val: int) -> int:
     coro.cache_clear()
 
     check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+
+async def test_cache_clear_pending_task() -> None:
+    @alru_cache()
+    async def coro() -> str:
+        await asyncio.sleep(0.5)
+        return "foo"
+
+    t = asyncio.create_task(coro())
+    await asyncio.sleep(0)
+    assert len(coro._LRUCacheWrapper__tasks) == 1  # type: ignore[attr-defined]
+    inner_task = next(iter(coro._LRUCacheWrapper__tasks))  # type: ignore[attr-defined]
+    assert not inner_task.done()
+
+    coro.cache_clear()
+    await inner_task
+
+    assert await t == "foo"
+    assert inner_task.done()
+
+
+async def test_cache_clear_ttl_callback(check_lru: Callable[..., None]) -> None:
+    @alru_cache(ttl=0.5)
+    async def coro() -> str:
+        return "foo"
+
+    await coro()
+    assert len(coro._LRUCacheWrapper__cache) == 1  # type: ignore[attr-defined]
+    cache_item = next(iter(coro._LRUCacheWrapper__cache.values()))  # type: ignore[attr-defined]
+    assert not cache_item.later_call.cancelled()
+
+    coro.cache_clear()
+
+    assert cache_item.later_call.cancelled()
+    await asyncio.sleep(0.5)