From e1a965cf8aed4a7d8bccc6d2fedc502570eb92d5 Mon Sep 17 00:00:00 2001
From: dkang-quora <dkang@quora.com>
Date: Thu, 18 Apr 2024 07:53:29 +0900
Subject: [PATCH] Fix leave_context_asyncio to handle cancelled asyncio task

---
 asynq/contexts.py   |  4 ++--
 asynq/decorators.py | 21 +++++++++++++--------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/asynq/contexts.py b/asynq/contexts.py
index 3e38871..828ad51 100644
--- a/asynq/contexts.py
+++ b/asynq/contexts.py
@@ -93,7 +93,7 @@ def leave_context_asyncio(context):
         debug.write("@async: -context: %s" % debug.str(context))
 
     task = asyncio.current_task()
-    del getattr(task, ASYNCIO_CONTEXT_FIELD)[id(context)]  # type: ignore
+    getattr(task, ASYNCIO_CONTEXT_FIELD, {}).pop(id(context), None)  # type: ignore
 
 
 def pause_contexts_asyncio(task):
@@ -104,7 +104,7 @@ def pause_contexts_asyncio(task):
 
 
 def resume_contexts_asyncio(task):
-    if not getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False):
+    if not getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, True):
         setattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, True)
         for ctx in getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values():
             ctx.resume()
diff --git a/asynq/decorators.py b/asynq/decorators.py
index 0c8d9a1..9178518 100644
--- a/asynq/decorators.py
+++ b/asynq/decorators.py
@@ -15,7 +15,7 @@
 
 import asyncio
 import inspect
-from typing import Any, Coroutine
+from typing import Any, Callable, Coroutine, Generic, TypeVar
 
 import qcore.decorators
 import qcore.helpers as core_helpers
@@ -27,6 +27,8 @@
 
 __traceback_hide__ = True
 
+_T = TypeVar("_T")
+
 
 def lazy(fn):
     """Converts a function into a lazy one - i.e. its call
@@ -147,10 +149,11 @@ def is_pure_async_fn(self):
         return True
 
 
-class PureAsyncDecorator(qcore.decorators.DecoratorBase):
+class PureAsyncDecorator(qcore.decorators.DecoratorBase, Generic[_T]):
     binder_cls = PureAsyncDecoratorBinder
+    fn: Callable[..., _T]
 
-    def __init__(self, fn, task_cls, kwargs={}, asyncio_fn=None):
+    def __init__(self, fn: Callable[..., _T], task_cls, kwargs={}, asyncio_fn=None):
         qcore.decorators.DecoratorBase.__init__(self, fn)
         self.task_cls = task_cls
         self.needs_wrapper = core_inspection.is_cython_or_generator(fn)
@@ -188,21 +191,23 @@ def _call_pure(self, args, kwargs):
         return self.task_cls(result, self.fn, args, kwargs, **self.kwargs)
 
 
-class AsyncDecoratorBinder(qcore.decorators.DecoratorBinder):
-    def asynq(self, *args, **kwargs):
+class AsyncDecoratorBinder(qcore.decorators.DecoratorBinder, Generic[_T]):
+    decorator: "AsyncDecorator[_T]"
+
+    def asynq(self, *args, **kwargs) -> async_task.AsyncTask[_T]:
         if self.instance is None:
             return self.decorator.asynq(*args, **kwargs)
         else:
             return self.decorator.asynq(self.instance, *args, **kwargs)
 
-    def asyncio(self, *args, **kwargs) -> Coroutine[Any, Any, Any]:
+    def asyncio(self, *args, **kwargs) -> Coroutine[Any, Any, _T]:
         if self.instance is None:
             return self.decorator.asyncio(*args, **kwargs)
         else:
             return self.decorator.asyncio(self.instance, *args, **kwargs)
 
 
-class AsyncDecorator(PureAsyncDecorator):
+class AsyncDecorator(PureAsyncDecorator[_T]):
     binder_cls = AsyncDecoratorBinder
 
     def __init__(self, fn, cls, kwargs={}, asyncio_fn=None):
@@ -211,7 +216,7 @@ def __init__(self, fn, cls, kwargs={}, asyncio_fn=None):
     def is_pure_async_fn(self):
         return False
 
-    def asynq(self, *args, **kwargs):
+    def asynq(self, *args: Any, **kwargs: Any) -> async_task.AsyncTask[_T]:
         return self._call_pure(args, kwargs)
 
     def name(self):