Skip to content

Commit

Permalink
Make AsyncContext no-op in asyncio mode
Browse files Browse the repository at this point in the history
  • Loading branch information
dkang-quora committed Apr 18, 2024
1 parent 90b9ca9 commit 4628dfc
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 54 deletions.
55 changes: 4 additions & 51 deletions asynq/contexts.py
Expand Up @@ -20,10 +20,6 @@
from ._debug import options as _debug_options


ASYNCIO_CONTEXT_FIELD = "_asynq_contexts"
ASYNCIO_CONTEXT_ACTIVE_FIELD = "_asynq_contexts_active"


class NonAsyncContext(object):
"""Indicates that context can't contain yield statements.
Expand All @@ -38,15 +34,11 @@ class NonAsyncContext(object):
"""

def __enter__(self):
if is_asyncio_mode():
enter_context_asyncio(self)
else:
if not is_asyncio_mode():
self._active_task = enter_context(self)

def __exit__(self, typ, val, tb):
if is_asyncio_mode():
leave_context_asyncio(self)
else:
if not is_asyncio_mode():
leave_context(self, self._active_task)

def pause(self):
Expand Down Expand Up @@ -75,40 +67,6 @@ def leave_context(context, active_task):
active_task._leave_context(context)


def enter_context_asyncio(context):
if _debug_options.DUMP_CONTEXTS:
debug.write("@async: +context: %s" % debug.str(context))

# since we are in asyncio mode, there is an active task
task = asyncio.current_task()

if hasattr(task, ASYNCIO_CONTEXT_FIELD):
getattr(task, ASYNCIO_CONTEXT_FIELD)[id(context)] = context
else:
setattr(task, ASYNCIO_CONTEXT_FIELD, {id(context): context})


def leave_context_asyncio(context):
if _debug_options.DUMP_CONTEXTS:
debug.write("@async: -context: %s" % debug.str(context))

task = asyncio.current_task()
getattr(task, ASYNCIO_CONTEXT_FIELD, {}).pop(id(context), None) # type: ignore


def pause_contexts_asyncio(task):
if getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False):
setattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False)
for ctx in reversed(list(getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values())):
ctx.pause()


def resume_contexts_asyncio(task):
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()


class AsyncContext(object):
"""Base class for contexts that should pause and resume during an async's function execution.
Expand All @@ -127,19 +85,14 @@ class AsyncContext(object):
"""

def __enter__(self):
if is_asyncio_mode():
enter_context_asyncio(self)
else:
if not is_asyncio_mode():
self._active_task = enter_context(self)

self.resume()
return self

def __exit__(self, ty, value, tb):
if is_asyncio_mode():
leave_context_asyncio(self)
self.pause()
else:
if not is_asyncio_mode():
leave_context(self, self._active_task)
self.pause()
del self._active_task
Expand Down
3 changes: 0 additions & 3 deletions asynq/decorators.py
Expand Up @@ -23,7 +23,6 @@

from . import async_task, futures
from .asynq_to_async import AsyncioMode, is_asyncio_mode, resolve_awaitables
from .contexts import pause_contexts_asyncio, resume_contexts_asyncio

__traceback_hide__ = True

Expand Down Expand Up @@ -114,7 +113,6 @@ async def wrapped(*_args, **_kwargs):

generator = fn(*_args, **_kwargs)
while True:
resume_contexts_asyncio(task)
try:
if exception is None:
result = generator.send(send)
Expand All @@ -125,7 +123,6 @@ async def wrapped(*_args, **_kwargs):
except StopIteration as exc:
return exc.value

pause_contexts_asyncio(task)
try:
send = await resolve_awaitables(result)
exception = None
Expand Down

0 comments on commit 4628dfc

Please sign in to comment.