|
3 | 3 | from __future__ import annotations |
4 | 4 |
|
5 | 5 | import asyncio |
6 | | -import sys |
7 | 6 | import typing as t |
8 | 7 | from collections.abc import Mapping |
9 | 8 | from contextvars import copy_context |
10 | | -from functools import partial, wraps |
| 9 | +from functools import wraps |
11 | 10 |
|
12 | 11 | if t.TYPE_CHECKING: |
13 | 12 | from collections.abc import Callable |
@@ -46,36 +45,13 @@ def _async_in_context( |
46 | 45 | ) -> Callable[..., t.Coroutine[T, U, V]]: |
47 | 46 | """ |
48 | 47 | Wrapper to run a coroutine in a persistent ContextVar Context. |
49 | | -
|
50 | | - Backports asyncio.create_task(context=...) behavior from Python 3.11 |
51 | 48 | """ |
52 | 49 | if context is None: |
53 | 50 | context = copy_context() |
54 | 51 |
|
55 | | - if sys.version_info >= (3, 11): |
56 | | - |
57 | | - @wraps(f) |
58 | | - async def run_in_context(*args, **kwargs): |
59 | | - coro = f(*args, **kwargs) |
60 | | - return await asyncio.create_task(coro, context=context) |
61 | | - |
62 | | - return run_in_context |
63 | | - |
64 | | - # don't need this backport when we require 3.11 |
65 | | - # context_holder so we have a modifiable container for later calls |
66 | | - context_holder = [context] # type: ignore[unreachable] |
67 | | - |
68 | | - async def preserve_context(f, *args, **kwargs): |
69 | | - """call a coroutine, preserving the context after it is called""" |
70 | | - try: |
71 | | - return await f(*args, **kwargs) |
72 | | - finally: |
73 | | - # persist changes to the context for future calls |
74 | | - context_holder[0] = copy_context() |
75 | | - |
76 | 52 | @wraps(f) |
77 | | - async def run_in_context_pre311(*args, **kwargs): |
78 | | - ctx = context_holder[0] |
79 | | - return await ctx.run(partial(asyncio.create_task, preserve_context(f, *args, **kwargs))) |
| 53 | + async def run_in_context(*args, **kwargs): |
| 54 | + coro = f(*args, **kwargs) |
| 55 | + return await asyncio.create_task(coro, context=context) |
80 | 56 |
|
81 | | - return run_in_context_pre311 |
| 57 | + return run_in_context |
0 commit comments