local.py

# Waiting for https://www.python.org/dev/peps/pep-0550/

import asyncio
from collections.abc import Coroutine

_original_task_factories = {}


def get_local():
    return getattr(asyncio.Task.current_task(), 'task_local', None)


class TaskWrapper(asyncio.Future):
    def __init__(self, coro, loop):
        self._passthru = False

        # Create task with wrapper which prevents passthru on late join
        this = self

        class CoroWrapper(Coroutine):
            def __await__(self):
                this._passthru = False
                yield from coro.__await__()

            def send(self, value):
                this._passthru = False
                return coro.send(value)

            def throw(self, typ, val=None, tb=None):
                this._passthru = False
                return coro.throw(typ, val, tb)

            def __getattr__(self, item):
                return getattr(coro, item)

        self._task = loop.create_task(CoroWrapper())
        self._task.task_local = {}
        if getattr(self._task, '_source_traceback', None):
            del getattr(self._task, '_source_traceback')[-3:]

        # Initialize Future object
        super().__init__(loop=loop)

        # Link cancel
        def check_cancel(fut):
            if fut.cancelled():
                self._task.cancel()
        self.add_done_callback(check_cancel)
        # Further `add_done_callback` may trigger local pass through
        self._passthru = True

        # Link done callback
        def task_done(task):
            if self.cancelled():
                return
            assert not self.done()
            if task.cancelled():
                self.cancel()
            else:
                exception = task.exception()
                if exception is not None:
                    self.set_exception(exception)
                else:
                    self.set_result(task.result())
        self._task.add_done_callback(task_done)

    def add_done_callback(self, *args, **kwargs):
        if self._passthru:
            self._passthru = False
            local = get_local()
            if local is not None:
                self._task.task_local = local
        return super().add_done_callback(*args, **kwargs)

    def with_local_reset(self):
        self._passthru = False
        return self

    def __getattr__(self, item):
        return getattr(self._task, item)


def reset_local(coro_or_future, *, loop=None):
    """Reset local to empty string within given routine.

    This works for:

    - newly created tasks
    - coroutines
    - awaitables

    For coroutines and awaitables, new tasks will be created with the give loop
    or current loop if not given.

    This doesn't work if:

    - task local is disabled for current/given loop
    - given task is already running
    - given futures are not instances of the internal `TaskWrapper` or subclass

    It is only the "reset local" part that is not working, the given future is
    returned untouched.
    """
    if isinstance(coro_or_future, TaskWrapper):
        return coro_or_future.with_local_reset()
    elif not asyncio.isfuture(coro_or_future):
        return reset_local(asyncio.ensure_future(coro_or_future, loop=loop))
    else:
        # we don't know how to reset local for random Future objects
        return coro_or_future


def task_factory(loop, coro):
    loop.set_task_factory(_original_task_factories.get(loop))
    try:
        return TaskWrapper(coro, loop)
    finally:
        loop.set_task_factory(task_factory)


def enable_task_local(loop=None):
    if loop is None:
        loop = asyncio.get_event_loop()
    if loop in _original_task_factories:
        return
    _original_task_factories[loop] = loop.get_task_factory()
    loop.set_task_factory(task_factory)


def disable_task_local(loop=None):
    if loop is None:
        loop = asyncio.get_event_loop()
    if loop in _original_task_factories:
        loop.set_task_factory(_original_task_factories.pop(loop))