Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/asgiref/sync.py: 21%

1import asyncio

2import asyncio.coroutines

3import contextvars

4import functools

5import inspect

6import os

7import sys

8import threading

9import warnings

10import weakref

11from collections.abc import Awaitable, Callable, Coroutine

12from concurrent.futures import Executor, Future, InvalidStateError, ThreadPoolExecutor

13from typing import (

14 TYPE_CHECKING,

15 Any,

16 Dict,

17 Generic,

18 List,

19 Optional,

20 ParamSpec,

21 TypeVar,

22 overload,

23)

25from .current_thread_executor import CurrentThreadExecutor

26from .local import Local, _rehome, _Storage

28if TYPE_CHECKING:

29 # This is not available to import at runtime

30 from _typeshed import OptExcInfo

32_F = TypeVar("_F", bound=Callable[..., Any])

33_P = ParamSpec("_P")

34_R = TypeVar("_R")

37def _restore_context(context: contextvars.Context) -> None:

38 # Check for changes in contextvars, and set them to the current

39 # context for downstream consumers

40 for cvar in context:

41 cvalue = context.get(cvar)

42 # asgiref is deliberately moving this context onto the current thread,

43 # so re-home any Local storage to it. This keeps Local data visible

44 # across async_to_sync / sync_to_async boundaries while leaving data

45 # merely inherited by an unrelated thread isolated (see asgiref.local).

46 if isinstance(cvalue, _Storage):

47 cvalue = _rehome(cvalue)

48 try:

49 if cvar.get() is not cvalue:

50 cvar.set(cvalue)

51 except LookupError:

52 cvar.set(cvalue)

55# Python 3.12 deprecates asyncio.iscoroutinefunction() as an alias for

56# inspect.iscoroutinefunction(), whilst also removing the _is_coroutine marker.

57# The latter is replaced with the inspect.markcoroutinefunction decorator.

58# Until 3.12 is the minimum supported Python version, provide a shim.

60if hasattr(inspect, "markcoroutinefunction"):

61 iscoroutinefunction = inspect.iscoroutinefunction

62 markcoroutinefunction: Callable[[_F], _F] = inspect.markcoroutinefunction

63else:

64 iscoroutinefunction = asyncio.iscoroutinefunction # type: ignore[assignment]

66 def markcoroutinefunction(func: _F) -> _F:

67 func._is_coroutine = asyncio.coroutines._is_coroutine # type: ignore

68 return func

71class AsyncSingleThreadContext:

72 """Context manager to run async code inside the same thread.

74 Normally, AsyncToSync functions run either inside a separate ThreadPoolExecutor or

75 the main event loop if it exists. This context manager ensures that all AsyncToSync

76 functions execute within the same thread.

78 This context manager is re-entrant, so only the outer-most call to

79 AsyncSingleThreadContext will set the context.

81 Usage:

83 >>> import asyncio

84 >>> with AsyncSingleThreadContext():

85 ... async_to_sync(asyncio.sleep(1))()

86 """

88 def __init__(self):

89 self.token = None

91 def __enter__(self):

92 try:

93 AsyncToSync.async_single_thread_context.get()

94 except LookupError:

95 self.token = AsyncToSync.async_single_thread_context.set(self)

97 return self

99 def __exit__(self, exc, value, tb):

100 if not self.token:

101 return

102

103 executor = AsyncToSync.context_to_thread_executor.pop(self, None)

104 if executor:

105 executor.shutdown()

106

107 AsyncToSync.async_single_thread_context.reset(self.token)

108

109

110class ThreadSensitiveContext:

111 """Async context manager to manage context for thread sensitive mode

112

113 This context manager controls which thread pool executor is used when in

114 thread sensitive mode. By default, a single thread pool executor is shared

115 within a process.

116

117 The ThreadSensitiveContext() context manager may be used to specify a

118 thread pool per context.

119

120 By default, only the outer-most ThreadSensitiveContext sets the context.

121 With an executor, this block uses that executor even inside another context

122 or an AsyncToSync call. On exit, the parent context is restored. The

123 executor must run one call at a time on one thread, for example

124 ThreadPoolExecutor(max_workers=1).

125

126 The caller owns the executor: it is not shut down on exit, so the executor

127 can be used again. A context with an executor can be entered only once.

128 Tasks created inside it must finish before it exits. Thread-sensitive calls

129 they make after the exit raise RuntimeError.

130

131 Usage:

132

133 >>> import time

134 >>> async with ThreadSensitiveContext():

135 ... await sync_to_async(time.sleep, 1)()

136 """

137

138 def __init__(self, *, executor: Executor | None = None) -> None:

139 self._executor = executor

140 self.token: contextvars.Token[ThreadSensitiveContext] | None = None

141 self._old_executor: CurrentThreadExecutor | None = None

142 # Set when a context with an executor exits. SyncToAsync then rejects

143 # calls from tasks that inherited the context and outlived it.

144 self._exited = False

145

146 async def __aenter__(self):

147 if self._executor is not None:

148 # An executor has been given. Use it in a new context rather than

149 # reusing an enclosing thread-sensitive context.

150 if self.token is not None:

151 raise RuntimeError("ThreadSensitiveContext is already entered")

152 if self._exited:

153 # Entering again would route tasks that outlived the first

154 # block to the second.

155 raise RuntimeError(

156 "ThreadSensitiveContext with an executor cannot be entered again"

157 )

158 self.token = SyncToAsync.thread_sensitive_context.set(self)

159 SyncToAsync.context_to_thread_executor[self] = self._executor

160 # An inherited CurrentThreadExecutor (via a parent AsyncToSync) has

161 # priority over this context in SyncToAsync.__call__. On entering,

162 # hide it to send calls in this block to the supplied executor. Any

163 # nested AsyncToSync call installs a new CurrentThreadExecutor, and

164 # will use that for further thread-sensitive work. (Nested calls

165 # sent to the supplied executor would deadlock: its only worker is

166 # waiting for the async code that's now waiting on the newly

167 # submitted code. See test_executor_context_across_bridges.)

168 self._old_executor = getattr(AsyncToSync.executors, "current", None)

169 AsyncToSync.executors.current = None

170 else:

171 try:

172 SyncToAsync.thread_sensitive_context.get()

173 except LookupError:

174 self.token = SyncToAsync.thread_sensitive_context.set(self)

175

176 return self

177

178 async def __aexit__(self, exc, value, tb):

179 if not self.token:

180 return

181

182 executor = SyncToAsync.context_to_thread_executor.pop(self, None)

183 SyncToAsync.thread_sensitive_context.reset(self.token)

184 self.token = None

185 if self._executor is not None:

186 AsyncToSync.executors.current = self._old_executor

187 self._old_executor = None

188 self._exited = True

189 # The caller owns the executor, so it is not shut down here.

190 return

191 if executor:

192 # The executor's worker thread may itself be waiting for this

193 # event loop, so a blocking shutdown() here would deadlock it.

194 # Join in a dedicated thread, not the loop's default executor:

195 # work queued there may itself be needed to unpark the worker,

196 # and joins occupying its slots would starve it.

197 future: "Future[None]" = Future()

198

199 def join() -> None:

200 executor.shutdown()

201 try:

202 future.set_result(None)

203 except InvalidStateError:

204 # The await below was cancelled while we were joining.

205 pass

206

207 threading.Thread(target=join, daemon=True).start()

208 await asyncio.wrap_future(future)

209

210

211class AsyncToSync(Generic[_P, _R]):

212 """

213 Utility class which turns an awaitable that only works on the thread with

214 the event loop into a synchronous callable that works in a subthread.

215

216 If the call stack contains an async loop, the code runs there.

217 Otherwise, the code runs in a new loop in a new thread.

218

219 Either way, this thread then pauses and waits to run any thread_sensitive

220 code called from further down the call stack using SyncToAsync, before

221 finally exiting once the async task returns.

222 """

223

224 # Keeps a reference to the CurrentThreadExecutor in local context, so that

225 # any sync_to_async inside the wrapped code can find it.

226 executors: "Local" = Local()

227

228 # When we can't find a CurrentThreadExecutor from the context, such as

229 # inside create_task, we'll look it up here from the running event loop.

230 loop_thread_executors: "Dict[asyncio.AbstractEventLoop, CurrentThreadExecutor]" = {}

231

232 async_single_thread_context: "contextvars.ContextVar[AsyncSingleThreadContext]" = (

233 contextvars.ContextVar("async_single_thread_context")

234 )

235

236 context_to_thread_executor: (

237 "weakref.WeakKeyDictionary[AsyncSingleThreadContext, ThreadPoolExecutor]"

238 ) = weakref.WeakKeyDictionary()

239

240 def __init__(

241 self,

242 awaitable: Callable[_P, Coroutine[Any, Any, _R]] | Callable[_P, Awaitable[_R]],

243 force_new_loop: bool = False,

244 ):

245 if not callable(awaitable) or (

246 not iscoroutinefunction(awaitable)

247 and not iscoroutinefunction(getattr(awaitable, "__call__", awaitable))

248 ):

249 # Python does not have very reliable detection of async functions

250 # (lots of false negatives) so this is just a warning.

251 warnings.warn(

252 "async_to_sync was passed a non-async-marked callable", stacklevel=2

253 )

254 self.awaitable = awaitable

255 try:

256 self.__self__ = self.awaitable.__self__ # type: ignore[union-attr]

257 except AttributeError:

258 pass

259 self.force_new_loop = force_new_loop

260

261 def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:

262 __traceback_hide__ = True # noqa: F841

263

264 main_event_loop = None

265 if not self.force_new_loop:

266 # There's no event loop in this thread. Look for the threadlocal if

267 # we're inside SyncToAsync

268 main_event_loop_pid = getattr(

269 SyncToAsync.threadlocal, "main_event_loop_pid", None

270 )

271 # We make sure the parent loop is from the same process - if

272 # they've forked, this is not going to be valid any more (#194)

273 if main_event_loop_pid and main_event_loop_pid == os.getpid():

274 main_event_loop = getattr(

275 SyncToAsync.threadlocal, "main_event_loop", None

276 )

277

278 # You can't call AsyncToSync from a thread with a running event loop

279 try:

280 asyncio.get_running_loop()

281 except RuntimeError:

282 pass

283 else:

284 raise RuntimeError(

285 "You cannot use AsyncToSync in the same thread as an async event loop - "

286 "just await the async function directly."

287 )

288

289 # Make a future for the return information

290 call_result: "Future[_R]" = Future()

291

292 # Make a CurrentThreadExecutor we'll use to idle in this thread - we

293 # need one for every sync frame, even if there's one above us in the

294 # same thread.

295 old_executor = getattr(self.executors, "current", None)

296 current_executor = CurrentThreadExecutor(old_executor)

297 self.executors.current = current_executor

298

299 # Wrapping context in list so it can be reassigned from within

300 # `main_wrap`.

301 context = [contextvars.copy_context()]

302

303 # Get task context so that parent task knows which task to propagate

304 # an asyncio.CancelledError to.

305 task_context = getattr(SyncToAsync.threadlocal, "task_context", None)

306

307 # Use call_soon_threadsafe to schedule a synchronous callback on the

308 # main event loop's thread if it's there, otherwise make a new loop

309 # in this thread.

310 try:

311 awaitable = self.main_wrap(

312 call_result,

313 sys.exc_info(),

314 task_context,

315 context,

316 # prepare an awaitable which can be passed as is to self.main_wrap,

317 # so that `args` and `kwargs` don't need to be

318 # destructured when passed to self.main_wrap

319 # (which is required by `ParamSpec`)

320 # as that may cause overlapping arguments

321 self.awaitable(*args, **kwargs),

322 )

323

324 async def new_loop_wrap() -> None:

325 loop = asyncio.get_running_loop()

326 self.loop_thread_executors[loop] = current_executor

327 try:

328 await awaitable

329 finally:

330 del self.loop_thread_executors[loop]

331

332 if main_event_loop is not None:

333 try:

334 main_event_loop.call_soon_threadsafe(

335 main_event_loop.create_task, awaitable

336 )

337 except RuntimeError:

338 running_in_main_event_loop = False

339 else:

340 running_in_main_event_loop = True

341 # Run the CurrentThreadExecutor until the future is done.

342 current_executor.run_until_future(call_result)

343 else:

344 running_in_main_event_loop = False

345

346 if not running_in_main_event_loop:

347 loop_executor = None

348

349 if self.async_single_thread_context.get(None):

350 single_thread_context = self.async_single_thread_context.get()

351

352 if single_thread_context in self.context_to_thread_executor:

353 loop_executor = self.context_to_thread_executor[

354 single_thread_context

355 ]

356 else:

357 loop_executor = ThreadPoolExecutor(max_workers=1)

358 self.context_to_thread_executor[single_thread_context] = (

359 loop_executor

360 )

361 else:

362 # Make our own event loop - in a new thread - and run inside that.

363 loop_executor = ThreadPoolExecutor(max_workers=1)

364

365 loop_future = loop_executor.submit(asyncio.run, new_loop_wrap())

366 # Run the CurrentThreadExecutor until the future is done.

367 current_executor.run_until_future(loop_future)

368 # Wait for future and/or allow for exception propagation

369 loop_future.result()

370 finally:

371 _restore_context(context[0])

372 # Restore old current thread executor state

373 self.executors.current = old_executor

374

375 # Wait for results from the future.

376 return call_result.result()

377

378 def __get__(self, parent: Any, objtype: Any) -> Callable[_P, _R]:

379 """

380 Include self for methods

381 """

382 func = functools.partial(self.__call__, parent)

383 return functools.update_wrapper(func, self.awaitable)

384

385 async def main_wrap(

386 self,

387 call_result: "Future[_R]",

388 exc_info: "OptExcInfo",

389 task_context: "Optional[List[asyncio.Task[Any]]]",

390 context: list[contextvars.Context],

391 awaitable: Coroutine[Any, Any, _R] | Awaitable[_R],

392 ) -> None:

393 """

394 Wraps the awaitable with something that puts the result into the

395 result/exception future.

396 """

397

398 __traceback_hide__ = True # noqa: F841

399

400 if context is not None:

401 _restore_context(context[0])

402

403 current_task = asyncio.current_task()

404 if current_task is not None and task_context is not None:

405 task_context.append(current_task)

406

407 try:

408 # If we have an exception, run the function inside the except block

409 # after raising it so exc_info is correctly populated.

410 if exc_info[1]:

411 try:

412 raise exc_info[1]

413 except BaseException:

414 result = await awaitable

415 else:

416 result = await awaitable

417 except BaseException as e:

418 call_result.set_exception(e)

419 else:

420 call_result.set_result(result)

421 finally:

422 if current_task is not None and task_context is not None:

423 task_context.remove(current_task)

424 context[0] = contextvars.copy_context()

425

426

427class SyncToAsync(Generic[_P, _R]):

428 """

429 Utility class which turns a synchronous callable into an awaitable that

430 runs in a threadpool. It also sets a threadlocal inside the thread so

431 calls to AsyncToSync can escape it.

432

433 If thread_sensitive is passed, the code will run in the same thread as any

434 outer code. This is needed for underlying Python code that is not

435 threadsafe (for example, code which handles SQLite database connections).

436

437 If the outermost program is async (i.e. SyncToAsync is outermost), then

438 this will be a dedicated single sub-thread that all sync code runs in,

439 one after the other. If the outermost program is sync (i.e. AsyncToSync is

440 outermost), this will just be the main thread. This is achieved by idling

441 with a CurrentThreadExecutor while AsyncToSync is blocking its sync parent,

442 rather than just blocking.

443

444 If executor is passed in, that will be used instead of the loop's default executor.

445 In order to pass in an executor, thread_sensitive must be set to False, otherwise

446 a TypeError will be raised.

447 """

448

449 # Storage for main event loop references

450 threadlocal = threading.local()

451

452 # Single-thread executor for thread-sensitive code

453 single_thread_executor = ThreadPoolExecutor(max_workers=1)

454

455 # Maintain a contextvar for the current execution context. Optionally used

456 # for thread sensitive mode.

457 thread_sensitive_context: "contextvars.ContextVar[ThreadSensitiveContext]" = (

458 contextvars.ContextVar("thread_sensitive_context")

459 )

460

461 # Contextvar that is used to detect if the single thread executor

462 # would be awaited on while already being used in the same context

463 deadlock_context: "contextvars.ContextVar[bool]" = contextvars.ContextVar(

464 "deadlock_context"

465 )

466

467 # Maintaining a weak reference to the context ensures that thread pools are

468 # erased once the context goes out of scope. This terminates the thread pool.

469 # A supplied executor is owned by the caller and is unaffected.

470 context_to_thread_executor: (

471 "weakref.WeakKeyDictionary[ThreadSensitiveContext, Executor]"

472 ) = weakref.WeakKeyDictionary()

473

474 def __init__(

475 self,

476 func: Callable[_P, _R],

477 thread_sensitive: bool = True,

478 executor: Optional["ThreadPoolExecutor"] = None,

479 context: contextvars.Context | None = None,

480 ) -> None:

481 if (

482 not callable(func)

483 or iscoroutinefunction(func)

484 or iscoroutinefunction(getattr(func, "__call__", func))

485 ):

486 raise TypeError("sync_to_async can only be applied to sync functions.")

487

488 functools.update_wrapper(self, func)

489 self.func = func

490 self.context = context

491

492 self._thread_sensitive = thread_sensitive

493 markcoroutinefunction(self)

494 if thread_sensitive and executor is not None:

495 raise TypeError("executor must not be set when thread_sensitive is True")

496 self._executor = executor

497 try:

498 self.__self__ = func.__self__ # type: ignore

499 except AttributeError:

500 pass

501

502 async def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:

503 __traceback_hide__ = True # noqa: F841

504 loop = asyncio.get_running_loop()

505

506 # Work out what thread to run the code in

507 if self._thread_sensitive:

508 current_thread_executor = getattr(AsyncToSync.executors, "current", None)

509 if current_thread_executor:

510 # If we have a parent sync thread above somewhere, use that

511 executor = current_thread_executor

512 elif self.thread_sensitive_context.get(None):

513 # If we have a way of retrieving the current context, attempt

514 # to use a per-context thread pool executor

515 thread_sensitive_context = self.thread_sensitive_context.get()

516

517 if thread_sensitive_context._exited:

518 # This task inherited a context with an executor and has

519 # outlived it. Creating an executor here would silently

520 # move the call to an unrelated thread.

521 raise RuntimeError(

522 "ThreadSensitiveContext has exited. Tasks created "

523 "inside a context with an executor must finish before "

524 "it exits."

525 )

526 if thread_sensitive_context in self.context_to_thread_executor:

527 # Re-use thread executor in current context

528 executor = self.context_to_thread_executor[thread_sensitive_context]

529 else:

530 # Create new thread executor in current context

531 executor = ThreadPoolExecutor(max_workers=1)

532 self.context_to_thread_executor[thread_sensitive_context] = executor

533 elif loop in AsyncToSync.loop_thread_executors:

534 # Re-use thread executor for running loop

535 executor = AsyncToSync.loop_thread_executors[loop]

536 elif self.deadlock_context.get(False):

537 raise RuntimeError(

538 "Single thread executor already being used, would deadlock"

539 )

540 else:

541 # Otherwise, we run it in a fixed single thread

542 executor = self.single_thread_executor

543 self.deadlock_context.set(True)

544 else:

545 # Use the passed in executor, or the loop's default if it is None

546 executor = self._executor

547

548 context = contextvars.copy_context() if self.context is None else self.context

549 # ``child`` is the deferred sync function to be run, with its args

550 # and kwargs bound.

551 child = functools.partial(self.func, *args, **kwargs)

552

553 # On the worker thread, thread_handler runs ``func(child)``. ``func``

554 # enters ``context`` (via context.run); then, inside it, ``run_child``

555 # re-homes any Local storage to the worker thread so it stays visible

556 # there (see _restore_context), and finally calls ``child``.

557 def func(child: Callable[[], _R]) -> _R:

558 def run_child() -> _R:

559 _restore_context(context)

560 return child()

561

562 return context.run(run_child)

563

564 task_context: list[asyncio.Task[Any]] = []

565

566 # Run the code in the right thread

567 exec_coro = loop.run_in_executor(

568 executor,

569 functools.partial(

570 self.thread_handler,

571 loop,

572 sys.exc_info(),

573 task_context,

574 func,

575 child,

576 ),

577 )

578 ret: _R

579 try:

580 ret = await asyncio.shield(exec_coro)

581 except asyncio.CancelledError:

582 cancel_parent = True

583 try:

584 task = task_context[0]

585 task.cancel()

586 try:

587 await task

588 cancel_parent = False

589 except asyncio.CancelledError:

590 pass

591 except IndexError:

592 pass

593 if exec_coro.done():

594 raise

595 if cancel_parent:

596 exec_coro.cancel()

597 ret = await exec_coro

598 finally:

599 if self.context is None:

600 _restore_context(context)

601 self.deadlock_context.set(False)

602

603 return ret

604

605 def __get__(

606 self, parent: Any, objtype: Any

607 ) -> Callable[_P, Coroutine[Any, Any, _R]]:

608 """

609 Include self for methods

610 """

611 func = functools.partial(self.__call__, parent)

612 return functools.update_wrapper(func, self.func)

613

614 def thread_handler(self, loop, exc_info, task_context, func, *args, **kwargs):

615 """

616 Wraps the sync application with exception handling.

617 """

618

619 __traceback_hide__ = True # noqa: F841

620

621 # Set the threadlocal for AsyncToSync

622 self.threadlocal.main_event_loop = loop

623 self.threadlocal.main_event_loop_pid = os.getpid()

624 self.threadlocal.task_context = task_context

625

626 # Run the function

627 # If we have an exception, run the function inside the except block

628 # after raising it so exc_info is correctly populated.

629 if exc_info[1]:

630 try:

631 raise exc_info[1]

632 except BaseException:

633 return func(*args, **kwargs)

634 else:

635 return func(*args, **kwargs)

636

637

638@overload

639def async_to_sync(

640 *,

641 force_new_loop: bool = False,

642) -> Callable[

643 [Callable[_P, Coroutine[Any, Any, _R]] | Callable[_P, Awaitable[_R]]],

644 Callable[_P, _R],

645]: ...

646

647

648@overload

649def async_to_sync(

650 awaitable: Callable[_P, Coroutine[Any, Any, _R]] | Callable[_P, Awaitable[_R]],

651 *,

652 force_new_loop: bool = False,

653) -> Callable[_P, _R]: ...

654

655

656def async_to_sync(

657 awaitable: None | (

658 Callable[_P, Coroutine[Any, Any, _R]] | Callable[_P, Awaitable[_R]]

659 ) = None,

660 *,

661 force_new_loop: bool = False,

662) -> (

663 Callable[

664 [Callable[_P, Coroutine[Any, Any, _R]] | Callable[_P, Awaitable[_R]]],

665 Callable[_P, _R],

666 ]

667 | Callable[_P, _R]

668):

669 if awaitable is None:

670 return lambda f: AsyncToSync(

671 f,

672 force_new_loop=force_new_loop,

673 )

674 return AsyncToSync(

675 awaitable,

676 force_new_loop=force_new_loop,

677 )

678

679

680@overload

681def sync_to_async(

682 *,

683 thread_sensitive: bool = True,

684 executor: Optional["ThreadPoolExecutor"] = None,

685 context: contextvars.Context | None = None,

686) -> Callable[[Callable[_P, _R]], Callable[_P, Coroutine[Any, Any, _R]]]: ...

687

688

689@overload

690def sync_to_async(

691 func: Callable[_P, _R],

692 *,

693 thread_sensitive: bool = True,

694 executor: Optional["ThreadPoolExecutor"] = None,

695 context: contextvars.Context | None = None,

696) -> Callable[_P, Coroutine[Any, Any, _R]]: ...

697

698

699def sync_to_async(

700 func: Callable[_P, _R] | None = None,

701 *,

702 thread_sensitive: bool = True,

703 executor: Optional["ThreadPoolExecutor"] = None,

704 context: contextvars.Context | None = None,

705) -> (

706 Callable[[Callable[_P, _R]], Callable[_P, Coroutine[Any, Any, _R]]]

707 | Callable[_P, Coroutine[Any, Any, _R]]

708):

709 if func is None:

710 return lambda f: SyncToAsync(

711 f,

712 thread_sensitive=thread_sensitive,

713 executor=executor,

714 context=context,

715 )

716 return SyncToAsync(

717 func,

718 thread_sensitive=thread_sensitive,

719 executor=executor,

720 context=context,

721 )