Python contextvars.copy_context() Examples

def run_sync(func: Callable[..., Any]) -> Callable[..., Coroutine[Any, None, None]]:
    """Ensure that the sync function is run within the event loop.

    If the *func* is not a coroutine it will be wrapped such that
    it runs in the default executor (use loop.set_default_executor
    to change). This ensures that synchronous functions do not
    block the event loop.
    """

    @wraps(func)
    async def _wrapper(*args: Any, **kwargs: Any) -> Any:
        loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(
            None, copy_context().run, partial(func, *args, **kwargs)
        )
        if isgenerator(result):
            return run_sync_iterable(result)  # type: ignore
        else:
            return result

    _wrapper._quart_async_wrapper = True  # type: ignore
    return _wrapper 

def threader():
    threads = []

    def start_fn(delay, fn):
        def thread_fn():
            time.sleep(delay)
            fn()

        target = functools.partial(contextvars.copy_context().run, thread_fn)
        thread = threading.Thread(target=target)
        thread.start()
        threads.append(thread)

    try:
        yield start_fn
    finally:
        for thread in threads:
            thread.join() 

def test_context_run_4(self):
        ctx1 = contextvars.Context()
        ctx2 = contextvars.Context()
        var = contextvars.ContextVar('var')

        def func2():
            self.assertIsNone(var.get(None))

        def func1():
            self.assertIsNone(var.get(None))
            var.set('spam')
            ctx2.run(func2)
            self.assertEqual(var.get(None), 'spam')

            cur = contextvars.copy_context()
            self.assertEqual(len(cur), 1)
            self.assertEqual(cur[var], 'spam')
            return cur

        returned_ctx = ctx1.run(func1)
        self.assertEqual(ctx1, returned_ctx)
        self.assertEqual(returned_ctx[var], 'spam')
        self.assertIn(var, returned_ctx) 

def __init__(self, coro, *, loop=None):
        super().__init__(loop=loop)
        if self._source_traceback:
            del self._source_traceback[-1]
        if not coroutines.iscoroutine(coro):
            # raise after Future.__init__(), attrs are required for __del__
            # prevent logging for pending task in __del__
            self._log_destroy_pending = False
            raise TypeError(f"a coroutine was expected, got {coro!r}")

        self._must_cancel = False
        self._fut_waiter = None
        self._coro = coro
        self._context = contextvars.copy_context()

        self._loop.call_soon(self.__step, context=self._context)
        _register_task(self) 

def __init__(self, coro, *, loop=None):
        super().__init__(loop=loop)
        if self._source_traceback:
            del self._source_traceback[-1]
        if not coroutines.iscoroutine(coro):
            # raise after Future.__init__(), attrs are required for __del__
            # prevent logging for pending task in __del__
            self._log_destroy_pending = False
            raise TypeError(f"a coroutine was expected, got {coro!r}")

        self._must_cancel = False
        self._fut_waiter = None
        self._coro = coro
        self._context = contextvars.copy_context()

        self._loop.call_soon(self.__step, context=self._context)
        _register_task(self) 

def __init__(self, coro, *, loop=None):
        super().__init__(loop=loop)
        if self._source_traceback:
            del self._source_traceback[-1]
        if not coroutines.iscoroutine(coro):
            # raise after Future.__init__(), attrs are required for __del__
            # prevent logging for pending task in __del__
            self._log_destroy_pending = False
            raise TypeError(f"a coroutine was expected, got {coro!r}")

        self._must_cancel = False
        self._fut_waiter = None
        self._coro = coro
        self._context = contextvars.copy_context()

        self._loop.call_soon(self.__step, context=self._context)
        _register_task(self) 

def __init__(self, coro, *, loop=None):
        super().__init__(loop=loop)
        if self._source_traceback:
            del self._source_traceback[-1]
        if not coroutines.iscoroutine(coro):
            # raise after Future.__init__(), attrs are required for __del__
            # prevent logging for pending task in __del__
            self._log_destroy_pending = False
            raise TypeError(f"a coroutine was expected, got {coro!r}")

        self._must_cancel = False
        self._fut_waiter = None
        self._coro = coro
        self._context = contextvars.copy_context()

        self._loop.call_soon(self.__step, context=self._context)
        _register_task(self) 

def submit_to_executor(func: Callable, *args, **kwargs) -> Future:
    """
    Submit a callable to the current context's executor pool and
    get back a future to monitor execution.

    Parameters
    ----------
    func : Callable
        Callable to be executed
    args
        Positional arguments to func
    kwargs
        Keyword arguments to func

    Returns
    -------
    Future

    """
    current_context = copy_context()
    return get_executor().submit(current_context.run, func, *args, **kwargs) 

def __call__(self, *args, **kwargs):
        # re-implementation of sync_to_async
        loop = asyncio.get_event_loop()
        context = contextvars.copy_context()
        child = functools.partial(self.func, *args, **kwargs)

        future = loop.run_in_executor(
            self.executor,
            functools.partial(
                self.thread_handler,
                loop,
                self.get_current_task(),
                sys.exc_info(),
                context.run,
                child,
            ),
        )
        return await asyncio.wait_for(future, timeout=None)


# The class is TitleCased, but we want to encourage use as a callable/decorator 

def inner(self, x):
        self.assertEqual(ctx_var.get(), x)
        await self.gen_inner(x)
        self.assertEqual(ctx_var.get(), x)

        # IOLoop.run_in_executor doesn't automatically copy context
        ctx = contextvars.copy_context()
        await self.io_loop.run_in_executor(None, lambda: ctx.run(self.thread_inner, x))
        self.assertEqual(ctx_var.get(), x)

        # Neither does asyncio's run_in_executor.
        await asyncio.get_event_loop().run_in_executor(
            None, lambda: ctx.run(self.thread_inner, x)
        )
        self.assertEqual(ctx_var.get(), x) 

def add_done_callback(self, fn, *, context=None):
        """Add a callback to be run when the future becomes done.

        The callback is called with a single argument - the future object. If
        the future is already done when this is called, the callback is
        scheduled with call_soon.
        """
        if self._state != _PENDING:
            self._loop.call_soon(fn, self, context=context)
        else:
            if context is None:
                context = contextvars.copy_context()
            self._callbacks.append((fn, context))

    # New method not in PEP 3148. 

def __init__(self, callback, args, loop, context=None):
        if context is None:
            context = contextvars.copy_context()
        self._context = context
        self._loop = loop
        self._callback = callback
        self._args = args
        self._cancelled = False
        self._repr = None
        if self._loop.get_debug():
            self._source_traceback = format_helpers.extract_stack(
                sys._getframe(1))
        else:
            self._source_traceback = None 

def __init__(self, callback, args, loop, context=None):
        if context is None:
            context = contextvars.copy_context()
        self._context = context
        self._loop = loop
        self._callback = callback
        self._args = args
        self._cancelled = False
        self._repr = None
        if self._loop.get_debug():
            self._source_traceback = format_helpers.extract_stack(
                sys._getframe(1))
        else:
            self._source_traceback = None 

def add_done_callback(self, fn, *, context=None):
        """Add a callback to be run when the future becomes done.

        The callback is called with a single argument - the future object. If
        the future is already done when this is called, the callback is
        scheduled with call_soon.
        """
        if self._state != _PENDING:
            self._loop.call_soon(fn, self, context=context)
        else:
            if context is None:
                context = contextvars.copy_context()
            self._callbacks.append((fn, context))

    # New method not in PEP 3148. 

def __init__(self, callback, args, loop, context=None):
        if context is None:
            context = contextvars.copy_context()
        self._context = context
        self._loop = loop
        self._callback = callback
        self._args = args
        self._cancelled = False
        self._repr = None
        if self._loop.get_debug():
            self._source_traceback = format_helpers.extract_stack(
                sys._getframe(1))
        else:
            self._source_traceback = None 

def test_context_get_context_1(self):
        ctx = contextvars.copy_context()
        self.assertIsInstance(ctx, contextvars.Context) 

def test_context_getset_5(self):
        c = contextvars.ContextVar('c', default=42)
        c.set([])

        def fun():
            c.set([])
            c.get().append(42)
            self.assertEqual(c.get(), [42])

        contextvars.copy_context().run(fun)
        self.assertEqual(c.get(), []) 

def __init__(self,
                 work,
                 args=(),
                 kw={},
                 max=100,
                 title=None,
                 tags=None,
                 comment='',
                 auto_save=None,
                 no_bar=False):
        title = work.__name__ if title is None else title
        self.ctx = contextvars.copy_context()
        self.parent = self.ctx.get(current_job, None)
        if auto_save is None:
            self.auto_save = True if self.parent is None else False
        else:
            self.auto_save = auto_save
        self.data = DataCollector(title, tags=tags, comment=comment, job=self)
        self.bar = ProgressBar(max=max, description=title, hiden=no_bar)
        self.out = widgets.Output()
        display(self.out)
        self.work_code = None
        code = compile(inspect.getsource(work),
                       f'defintion of {work.__name__}', 'single')
        for c in code.co_consts:
            if isinstance(c, type(code)) and c.co_name == work.__name__:
                self.work_code = c
                break
        self.work = work
        self.args = args
        self.kw = kw 

def add_done_callback(self, fn, *, context=None):
        """Add a callback to be run when the future becomes done.

        The callback is called with a single argument - the future object. If
        the future is already done when this is called, the callback is
        scheduled with call_soon.
        """
        if self._state != _PENDING:
            self._loop.call_soon(fn, self, context=context)
        else:
            if context is None:
                context = contextvars.copy_context()
            self._callbacks.append((fn, context))

    # New method not in PEP 3148. 

def run_in_executor_with_context(
    func: Callable[..., _T], *args: Any, loop: Optional[AbstractEventLoop] = None
) -> Awaitable[_T]:
    """
    Run a function in an executor, but make sure it uses the same contextvars.
    This is required so that the function will see the right application.

    See also: https://bugs.python.org/issue34014
    """
    loop = loop or get_event_loop()
    ctx: contextvars.Context = contextvars.copy_context()

    return loop.run_in_executor(None, ctx.run, func, *args) 

def schedule(self, x, context_map=None):
        """Schedule a task."""
        if context_map:
            ctx = copy_context()
            ctx.run(lambda: [k.set(v) for k, v in context_map.items()])
            fut = ctx.run(asyncio.ensure_future, x, loop=self)
        else:
            fut = asyncio.ensure_future(x, loop=self)
        self._tasks.append(fut)
        return fut 

def init_stack(self, generator):
        """Create a new stack for the given generator."""
        self._contexts[generator] = copy_context() 

def __init__(self, *, loop, source, repeat, callback, args, context=None):
        super().__init__(callback, args, loop)

        if sys.version_info[:2] >= (3, 7) and context is None:
            import contextvars
            context = contextvars.copy_context()
        self._context = context
        self._source = source
        self._repeat = repeat
        loop._handlers.add(self)
        source.set_callback(self.__callback__, self)
        source.attach(loop._context) 

def __init__(self, callback, args, loop, context=None):
        if context is None:
            context = contextvars.copy_context()
        self._context = context
        self._loop = loop
        self._callback = callback
        self._args = args
        self._cancelled = False
        self._repr = None
        if self._loop.get_debug():
            self._source_traceback = format_helpers.extract_stack(
                sys._getframe(1))
        else:
            self._source_traceback = None 

def add_done_callback(self, fn, *, context=None):
        """Add a callback to be run when the future becomes done.

        The callback is called with a single argument - the future object. If
        the future is already done when this is called, the callback is
        scheduled with call_soon.
        """
        if self._state != _PENDING:
            self._loop.call_soon(fn, self, context=context)
        else:
            if context is None:
                context = contextvars.copy_context()
            self._callbacks.append((fn, context))

    # New method not in PEP 3148. 

def inner(self, x):
        self.assertEqual(ctx_var.get(), x)
        await self.gen_inner(x)
        self.assertEqual(ctx_var.get(), x)

        # IOLoop.run_in_executor doesn't automatically copy context
        ctx = contextvars.copy_context()
        await self.io_loop.run_in_executor(None, lambda: ctx.run(self.thread_inner, x))
        self.assertEqual(ctx_var.get(), x)

        # Neither does asyncio's run_in_executor.
        await asyncio.get_event_loop().run_in_executor(
            None, lambda: ctx.run(self.thread_inner, x)
        )
        self.assertEqual(ctx_var.get(), x) 

def run_sync_iterable(iterable: Generator[Any, None, None]) -> AsyncGenerator[Any, None]:
    async def _gen_wrapper() -> AsyncGenerator[Any, None]:
        # Wrap the generator such that each iteration runs
        # in the executor. Then rationalise the raised
        # errors so that it ends.
        def _inner() -> Any:
            # https://bugs.python.org/issue26221
            # StopIteration errors are swallowed by the
            # run_in_exector method
            try:
                return next(iterable)
            except StopIteration:
                raise StopAsyncIteration()

        loop = asyncio.get_running_loop()
        while True:
            try:
                yield await loop.run_in_executor(None, copy_context().run, _inner)
            except StopAsyncIteration:
                return

    return _gen_wrapper() 

def invoke(
        fn: callbacks.BaseFn,
        *args: Any,
        settings: Optional[configuration.OperatorSettings] = None,
        cause: Optional[causation.BaseCause] = None,
        **kwargs: Any,
) -> Any:
    """
    Invoke a single function, but safely for the main asyncio process.

    Used mostly for handler functions, and potentially slow & blocking code.
    Other callbacks are called directly, and are expected to be synchronous
    (such as handler-selecting (lifecycles) and resource-filtering (``when=``)).

    A full set of the arguments is provided, expanding the cause to some easily
    usable aliases. The function is expected to accept ``**kwargs`` for the args
    that it does not use -- for forward compatibility with the new features.

    The synchronous methods are executed in the executor (threads or processes),
    thus making it non-blocking for the main event loop of the operator.
    See: https://pymotw.com/3/asyncio/executors.html
    """
    if is_async_fn(fn):
        kwargs = build_kwargs(cause=cause, _sync=False, **kwargs)
        result = await fn(*args, **kwargs)  # type: ignore
    else:
        kwargs = build_kwargs(cause=cause, _sync=True, **kwargs)

        # Not that we want to use functools, but for executors kwargs, it is officially recommended:
        # https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.run_in_executor
        real_fn = functools.partial(fn, *args, **kwargs)

        # Copy the asyncio context from current thread to the handlr's thread.
        # It can be copied 2+ times if there are sub-sub-handlers (rare case).
        context = contextvars.copy_context()
        real_fn = functools.partial(context.run, real_fn)

        # Prevent orphaned threads during daemon/handler cancellation. It is better to be stuck
        # in the task than to have orphan threads which deplete the executor's pool capacity.
        # Cancellation is postponed until the thread exits, but it happens anyway (for consistency).
        # Note: the docs say the result is a future, but typesheds say it is a coroutine => cast()!
        loop = asyncio.get_event_loop()
        executor = settings.execution.executor if settings is not None else None
        future = cast(asyncio_Future, loop.run_in_executor(executor, real_fn))
        cancellation: Optional[asyncio.CancelledError] = None
        while not future.done():
            try:
                await asyncio.shield(future)  # slightly expensive: creates tasks
            except asyncio.CancelledError as e:
                cancellation = e
        if cancellation is not None:
            raise cancellation
        result = future.result()

    return result 

def test_coroutine_middleware():
    before_call_trace.clear()
    create_call_trace.clear()
    run_call_trace_pre.clear()
    run_call_trace_post.clear()

    var = contextvars.ContextVar('caller_var')  # type: contextvars.ContextVar[str]
    var.set('yes man')

    test_context = {
        'caller_var': None,
        'middleware_var': None
    }  # type: Dict[str, Any]

    # noinspection PyCompatibility
    async def coroutine():
        run_call_trace_pre.append('target')

        test_context['caller_var'] = var.get('default_cv')
        for context_var in contextvars.copy_context().keys():
            if context_var.name == 'middleware_var':
                test_context['middleware_var'] = context_var.get('default_mv')
        await asyncio.sleep(0.05)

        run_call_trace_post.append('target')

        return 'foo_coroutine_returned_this'

    thread = AsyncEventLoopThread([
        SpecialCoroutineMiddleware(42),
        TracingCoroutineMiddleware(),
    ])
    thread.start()

    future = thread.run_coroutine(coroutine())

    await asyncio.sleep(0.01)

    assert future.result() == 'foo_coroutine_returned_this'

    thread.join()

    assert before_call_trace == ['SpecialCoroutineMiddleware', 'TracingCoroutineMiddleware']
    assert create_call_trace == ['TracingCoroutineMiddleware', 'SpecialCoroutineMiddleware']
    assert run_call_trace_pre == ['SpecialCoroutineMiddleware', 'TracingCoroutineMiddleware', 'target']
    assert run_call_trace_post == ['target', 'TracingCoroutineMiddleware', 'SpecialCoroutineMiddleware']

    assert test_context['caller_var'] == 'yes man'
    assert test_context['middleware_var'] == 42


# noinspection PyCompatibility