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

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 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() != 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 This context manager is re-entrant, so only the outer-most call to

121 ThreadSensitiveContext will set the context.

122

123 Usage:

124

125 >>> import time

126 >>> async with ThreadSensitiveContext():

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

128 """

129

130 def __init__(self):

131 self.token = None

132

133 async def __aenter__(self):

134 try:

135 SyncToAsync.thread_sensitive_context.get()

136 except LookupError:

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

138

139 return self

140

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

142 if not self.token:

143 return

144

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

146 SyncToAsync.thread_sensitive_context.reset(self.token)

147 if executor:

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

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

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

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

152 # and joins occupying its slots would starve it.

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

154

155 def join() -> None:

156 executor.shutdown()

157 try:

158 future.set_result(None)

159 except InvalidStateError:

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

161 pass

162

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

164 await asyncio.wrap_future(future)

165

166

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

168 """

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

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

171

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

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

174

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

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

177 finally exiting once the async task returns.

178 """

179

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

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

182 executors: "Local" = Local()

183

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

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

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

187

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

189 contextvars.ContextVar("async_single_thread_context")

190 )

191

192 context_to_thread_executor: (

193 "weakref.WeakKeyDictionary[AsyncSingleThreadContext, ThreadPoolExecutor]"

194 ) = weakref.WeakKeyDictionary()

195

196 def __init__(

197 self,

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

199 force_new_loop: bool = False,

200 ):

201 if not callable(awaitable) or (

202 not iscoroutinefunction(awaitable)

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

204 ):

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

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

207 warnings.warn(

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

209 )

210 self.awaitable = awaitable

211 try:

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

213 except AttributeError:

214 pass

215 self.force_new_loop = force_new_loop

216

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

218 __traceback_hide__ = True # noqa: F841

219

220 main_event_loop = None

221 if not self.force_new_loop:

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

223 # we're inside SyncToAsync

224 main_event_loop_pid = getattr(

225 SyncToAsync.threadlocal, "main_event_loop_pid", None

226 )

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

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

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

230 main_event_loop = getattr(

231 SyncToAsync.threadlocal, "main_event_loop", None

232 )

233

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

235 try:

236 asyncio.get_running_loop()

237 except RuntimeError:

238 pass

239 else:

240 raise RuntimeError(

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

242 "just await the async function directly."

243 )

244

245 # Make a future for the return information

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

247

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

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

250 # same thread.

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

252 current_executor = CurrentThreadExecutor(old_executor)

253 self.executors.current = current_executor

254

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

256 # `main_wrap`.

257 context = [contextvars.copy_context()]

258

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

260 # an asyncio.CancelledError to.

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

262

263 # Use call_soon_threadsafe to schedule a synchronous callback on the

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

265 # in this thread.

266 try:

267 awaitable = self.main_wrap(

268 call_result,

269 sys.exc_info(),

270 task_context,

271 context,

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

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

274 # destructured when passed to self.main_wrap

275 # (which is required by `ParamSpec`)

276 # as that may cause overlapping arguments

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

278 )

279

280 async def new_loop_wrap() -> None:

281 loop = asyncio.get_running_loop()

282 self.loop_thread_executors[loop] = current_executor

283 try:

284 await awaitable

285 finally:

286 del self.loop_thread_executors[loop]

287

288 if main_event_loop is not None:

289 try:

290 main_event_loop.call_soon_threadsafe(

291 main_event_loop.create_task, awaitable

292 )

293 except RuntimeError:

294 running_in_main_event_loop = False

295 else:

296 running_in_main_event_loop = True

297 # Run the CurrentThreadExecutor until the future is done.

298 current_executor.run_until_future(call_result)

299 else:

300 running_in_main_event_loop = False

301

302 if not running_in_main_event_loop:

303 loop_executor = None

304

305 if self.async_single_thread_context.get(None):

306 single_thread_context = self.async_single_thread_context.get()

307

308 if single_thread_context in self.context_to_thread_executor:

309 loop_executor = self.context_to_thread_executor[

310 single_thread_context

311 ]

312 else:

313 loop_executor = ThreadPoolExecutor(max_workers=1)

314 self.context_to_thread_executor[single_thread_context] = (

315 loop_executor

316 )

317 else:

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

319 loop_executor = ThreadPoolExecutor(max_workers=1)

320

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

322 # Run the CurrentThreadExecutor until the future is done.

323 current_executor.run_until_future(loop_future)

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

325 loop_future.result()

326 finally:

327 _restore_context(context[0])

328 # Restore old current thread executor state

329 self.executors.current = old_executor

330

331 # Wait for results from the future.

332 return call_result.result()

333

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

335 """

336 Include self for methods

337 """

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

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

340

341 async def main_wrap(

342 self,

343 call_result: "Future[_R]",

344 exc_info: "OptExcInfo",

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

346 context: list[contextvars.Context],

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

348 ) -> None:

349 """

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

351 result/exception future.

352 """

353

354 __traceback_hide__ = True # noqa: F841

355

356 if context is not None:

357 _restore_context(context[0])

358

359 current_task = asyncio.current_task()

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

361 task_context.append(current_task)

362

363 try:

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

365 # after raising it so exc_info is correctly populated.

366 if exc_info[1]:

367 try:

368 raise exc_info[1]

369 except BaseException:

370 result = await awaitable

371 else:

372 result = await awaitable

373 except BaseException as e:

374 call_result.set_exception(e)

375 else:

376 call_result.set_result(result)

377 finally:

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

379 task_context.remove(current_task)

380 context[0] = contextvars.copy_context()

381

382

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

384 """

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

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

387 calls to AsyncToSync can escape it.

388

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

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

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

392

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

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

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

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

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

398 rather than just blocking.

399

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

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

402 a TypeError will be raised.

403 """

404

405 # Storage for main event loop references

406 threadlocal = threading.local()

407

408 # Single-thread executor for thread-sensitive code

409 single_thread_executor = ThreadPoolExecutor(max_workers=1)

410

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

412 # for thread sensitive mode.

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

414 contextvars.ContextVar("thread_sensitive_context")

415 )

416

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

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

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

420 "deadlock_context"

421 )

422

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

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

425 context_to_thread_executor: (

426 "weakref.WeakKeyDictionary[ThreadSensitiveContext, ThreadPoolExecutor]"

427 ) = weakref.WeakKeyDictionary()

428

429 def __init__(

430 self,

431 func: Callable[_P, _R],

432 thread_sensitive: bool = True,

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

434 context: contextvars.Context | None = None,

435 ) -> None:

436 if (

437 not callable(func)

438 or iscoroutinefunction(func)

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

440 ):

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

442

443 functools.update_wrapper(self, func)

444 self.func = func

445 self.context = context

446

447 self._thread_sensitive = thread_sensitive

448 markcoroutinefunction(self)

449 if thread_sensitive and executor is not None:

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

451 self._executor = executor

452 try:

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

454 except AttributeError:

455 pass

456

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

458 __traceback_hide__ = True # noqa: F841

459 loop = asyncio.get_running_loop()

460

461 # Work out what thread to run the code in

462 if self._thread_sensitive:

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

464 if current_thread_executor:

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

466 executor = current_thread_executor

467 elif self.thread_sensitive_context.get(None):

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

469 # to use a per-context thread pool executor

470 thread_sensitive_context = self.thread_sensitive_context.get()

471

472 if thread_sensitive_context in self.context_to_thread_executor:

473 # Re-use thread executor in current context

474 executor = self.context_to_thread_executor[thread_sensitive_context]

475 else:

476 # Create new thread executor in current context

477 executor = ThreadPoolExecutor(max_workers=1)

478 self.context_to_thread_executor[thread_sensitive_context] = executor

479 elif loop in AsyncToSync.loop_thread_executors:

480 # Re-use thread executor for running loop

481 executor = AsyncToSync.loop_thread_executors[loop]

482 elif self.deadlock_context.get(False):

483 raise RuntimeError(

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

485 )

486 else:

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

488 executor = self.single_thread_executor

489 self.deadlock_context.set(True)

490 else:

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

492 executor = self._executor

493

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

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

496 # and kwargs bound.

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

498

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

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

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

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

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

504 def run_child() -> _R:

505 _restore_context(context)

506 return child()

507

508 return context.run(run_child)

509

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

511

512 # Run the code in the right thread

513 exec_coro = loop.run_in_executor(

514 executor,

515 functools.partial(

516 self.thread_handler,

517 loop,

518 sys.exc_info(),

519 task_context,

520 func,

521 child,

522 ),

523 )

524 ret: _R

525 try:

526 ret = await asyncio.shield(exec_coro)

527 except asyncio.CancelledError:

528 cancel_parent = True

529 try:

530 task = task_context[0]

531 task.cancel()

532 try:

533 await task

534 cancel_parent = False

535 except asyncio.CancelledError:

536 pass

537 except IndexError:

538 pass

539 if exec_coro.done():

540 raise

541 if cancel_parent:

542 exec_coro.cancel()

543 ret = await exec_coro

544 finally:

545 if self.context is None:

546 _restore_context(context)

547 self.deadlock_context.set(False)

548

549 return ret

550

551 def __get__(

552 self, parent: Any, objtype: Any

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

554 """

555 Include self for methods

556 """

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

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

559

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

561 """

562 Wraps the sync application with exception handling.

563 """

564

565 __traceback_hide__ = True # noqa: F841

566

567 # Set the threadlocal for AsyncToSync

568 self.threadlocal.main_event_loop = loop

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

570 self.threadlocal.task_context = task_context

571

572 # Run the function

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

574 # after raising it so exc_info is correctly populated.

575 if exc_info[1]:

576 try:

577 raise exc_info[1]

578 except BaseException:

579 return func(*args, **kwargs)

580 else:

581 return func(*args, **kwargs)

582

583

584@overload

585def async_to_sync(

586 *,

587 force_new_loop: bool = False,

588) -> Callable[

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

590 Callable[_P, _R],

591]: ...

592

593

594@overload

595def async_to_sync(

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

597 *,

598 force_new_loop: bool = False,

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

600

601

602def async_to_sync(

603 awaitable: None | (

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

605 ) = None,

606 *,

607 force_new_loop: bool = False,

608) -> (

609 Callable[

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

611 Callable[_P, _R],

612 ]

613 | Callable[_P, _R]

614):

615 if awaitable is None:

616 return lambda f: AsyncToSync(

617 f,

618 force_new_loop=force_new_loop,

619 )

620 return AsyncToSync(

621 awaitable,

622 force_new_loop=force_new_loop,

623 )

624

625

626@overload

627def sync_to_async(

628 *,

629 thread_sensitive: bool = True,

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

631 context: contextvars.Context | None = None,

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

633

634

635@overload

636def sync_to_async(

637 func: Callable[_P, _R],

638 *,

639 thread_sensitive: bool = True,

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

641 context: contextvars.Context | None = None,

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

643

644

645def sync_to_async(

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

647 *,

648 thread_sensitive: bool = True,

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

650 context: contextvars.Context | None = None,

651) -> (

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

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

654):

655 if func is None:

656 return lambda f: SyncToAsync(

657 f,

658 thread_sensitive=thread_sensitive,

659 executor=executor,

660 context=context,

661 )

662 return SyncToAsync(

663 func,

664 thread_sensitive=thread_sensitive,

665 executor=executor,

666 context=context,

667 )