_asyncio.py 105 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201
  1. from __future__ import annotations
  2. import array
  3. import asyncio
  4. import concurrent.futures
  5. import math
  6. import os
  7. import socket
  8. import sys
  9. import threading
  10. import weakref
  11. from asyncio import (
  12. AbstractEventLoop,
  13. CancelledError,
  14. all_tasks,
  15. create_task,
  16. current_task,
  17. get_running_loop,
  18. sleep,
  19. )
  20. from asyncio.base_events import _run_until_complete_cb # type: ignore[attr-defined]
  21. from collections import OrderedDict, deque
  22. from collections.abc import (
  23. AsyncGenerator,
  24. AsyncIterator,
  25. Awaitable,
  26. Callable,
  27. Collection,
  28. Coroutine,
  29. Iterable,
  30. Sequence,
  31. )
  32. from concurrent.futures import Future
  33. from contextlib import AbstractContextManager, suppress
  34. from contextvars import Context, copy_context
  35. from dataclasses import dataclass, field
  36. from functools import partial, wraps
  37. from inspect import (
  38. CORO_RUNNING,
  39. CORO_SUSPENDED,
  40. getcoroutinestate,
  41. )
  42. from io import IOBase
  43. from os import PathLike
  44. from queue import Queue
  45. from signal import Signals
  46. from socket import AddressFamily, SocketKind
  47. from threading import Thread
  48. from types import CodeType, TracebackType
  49. from typing import (
  50. IO,
  51. TYPE_CHECKING,
  52. Any,
  53. Literal,
  54. ParamSpec,
  55. TypeVar,
  56. cast,
  57. )
  58. from weakref import WeakKeyDictionary
  59. from .. import (
  60. CapacityLimiterStatistics,
  61. EventStatistics,
  62. LockStatistics,
  63. TaskInfo,
  64. abc,
  65. )
  66. from .._core._eventloop import (
  67. claim_worker_thread,
  68. set_current_async_library,
  69. threadlocals,
  70. )
  71. from .._core._exceptions import (
  72. BrokenResourceError,
  73. BusyResourceError,
  74. ClosedResourceError,
  75. EndOfStream,
  76. RunFinishedError,
  77. WouldBlock,
  78. )
  79. from .._core._sockets import convert_ipv6_sockaddr
  80. from .._core._streams import create_memory_object_stream
  81. from .._core._synchronization import (
  82. CapacityLimiter as BaseCapacityLimiter,
  83. )
  84. from .._core._synchronization import Event as BaseEvent
  85. from .._core._synchronization import Lock as BaseLock
  86. from .._core._synchronization import (
  87. ResourceGuard,
  88. SemaphoreStatistics,
  89. )
  90. from .._core._synchronization import Semaphore as BaseSemaphore
  91. from .._core._tasks import CancelScope as BaseCancelScope
  92. from .._core._tasks import TaskHandle
  93. from ..abc import (
  94. AsyncBackend,
  95. IPSockAddrType,
  96. SocketListener,
  97. UDPPacketType,
  98. UNIXDatagramPacketType,
  99. )
  100. from ..abc._tasks import call_for_coroutine, get_callable_name, get_coro_name
  101. from ..lowlevel import RunVar, _run_vars
  102. if TYPE_CHECKING:
  103. from _typeshed import FileDescriptorLike
  104. from ..abc._eventloop import StrOrBytesPath
  105. from ..streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
  106. else:
  107. FileDescriptorLike = object
  108. if sys.version_info >= (3, 11):
  109. from asyncio import Runner
  110. from typing import Self, TypeVarTuple, Unpack
  111. else:
  112. import contextvars
  113. import enum
  114. import signal
  115. from asyncio import coroutines, events, exceptions, tasks
  116. from exceptiongroup import BaseExceptionGroup
  117. from typing_extensions import Self, TypeVarTuple, Unpack
  118. class _State(enum.Enum):
  119. CREATED = "created"
  120. INITIALIZED = "initialized"
  121. CLOSED = "closed"
  122. class Runner:
  123. # Copied from CPython 3.11
  124. def __init__(
  125. self,
  126. *,
  127. debug: bool | None = None,
  128. loop_factory: Callable[[], AbstractEventLoop] | None = None,
  129. ):
  130. self._state = _State.CREATED
  131. self._debug = debug
  132. self._loop_factory = loop_factory
  133. self._loop: AbstractEventLoop | None = None
  134. self._context = None
  135. self._interrupt_count = 0
  136. self._set_event_loop = False
  137. def __enter__(self) -> Self:
  138. self._lazy_init()
  139. return self
  140. def __exit__(
  141. self,
  142. exc_type: type[BaseException] | None,
  143. exc_val: BaseException | None,
  144. exc_tb: TracebackType | None,
  145. ) -> None:
  146. self.close()
  147. def close(self) -> None:
  148. """Shutdown and close event loop."""
  149. loop = self._loop
  150. if self._state is not _State.INITIALIZED or loop is None:
  151. return
  152. try:
  153. _cancel_all_tasks(loop)
  154. loop.run_until_complete(loop.shutdown_asyncgens())
  155. if hasattr(loop, "shutdown_default_executor"):
  156. loop.run_until_complete(loop.shutdown_default_executor())
  157. else:
  158. loop.run_until_complete(_shutdown_default_executor(loop))
  159. finally:
  160. if self._set_event_loop:
  161. events.set_event_loop(None)
  162. loop.close()
  163. self._loop = None
  164. self._state = _State.CLOSED
  165. def get_loop(self) -> AbstractEventLoop:
  166. """Return embedded event loop."""
  167. self._lazy_init()
  168. return self._loop
  169. def run(self, coro: Coroutine[T_Retval], *, context=None) -> T_Retval:
  170. """Run a coroutine inside the embedded event loop."""
  171. if not coroutines.iscoroutine(coro):
  172. raise ValueError(f"a coroutine was expected, got {coro!r}")
  173. if events._get_running_loop() is not None:
  174. # fail fast with short traceback
  175. raise RuntimeError(
  176. "Runner.run() cannot be called from a running event loop"
  177. )
  178. self._lazy_init()
  179. if context is None:
  180. context = self._context
  181. task = context.run(self._loop.create_task, coro)
  182. if (
  183. threading.current_thread() is threading.main_thread()
  184. and signal.getsignal(signal.SIGINT) is signal.default_int_handler
  185. ):
  186. sigint_handler = partial(self._on_sigint, main_task=task)
  187. try:
  188. signal.signal(signal.SIGINT, sigint_handler)
  189. except ValueError:
  190. # `signal.signal` may throw if `threading.main_thread` does
  191. # not support signals (e.g. embedded interpreter with signals
  192. # not registered - see gh-91880)
  193. sigint_handler = None
  194. else:
  195. sigint_handler = None
  196. self._interrupt_count = 0
  197. try:
  198. return self._loop.run_until_complete(task)
  199. except exceptions.CancelledError:
  200. if self._interrupt_count > 0:
  201. uncancel = getattr(task, "uncancel", None)
  202. if uncancel is not None and uncancel() == 0:
  203. raise KeyboardInterrupt # noqa: B904
  204. raise # CancelledError
  205. finally:
  206. if (
  207. sigint_handler is not None
  208. and signal.getsignal(signal.SIGINT) is sigint_handler
  209. ):
  210. signal.signal(signal.SIGINT, signal.default_int_handler)
  211. def _lazy_init(self) -> None:
  212. if self._state is _State.CLOSED:
  213. raise RuntimeError("Runner is closed")
  214. if self._state is _State.INITIALIZED:
  215. return
  216. if self._loop_factory is None:
  217. self._loop = events.new_event_loop()
  218. if not self._set_event_loop:
  219. # Call set_event_loop only once to avoid calling
  220. # attach_loop multiple times on child watchers
  221. events.set_event_loop(self._loop)
  222. self._set_event_loop = True
  223. else:
  224. self._loop = self._loop_factory()
  225. if self._debug is not None:
  226. self._loop.set_debug(self._debug)
  227. self._context = contextvars.copy_context()
  228. self._state = _State.INITIALIZED
  229. def _on_sigint(self, signum, frame, main_task: asyncio.Task) -> None:
  230. self._interrupt_count += 1
  231. if self._interrupt_count == 1 and not main_task.done():
  232. main_task.cancel()
  233. # wakeup loop if it is blocked by select() with long timeout
  234. self._loop.call_soon_threadsafe(lambda: None)
  235. return
  236. raise KeyboardInterrupt()
  237. def _cancel_all_tasks(loop: AbstractEventLoop) -> None:
  238. to_cancel = tasks.all_tasks(loop)
  239. if not to_cancel:
  240. return
  241. for task in to_cancel:
  242. task.cancel()
  243. loop.run_until_complete(tasks.gather(*to_cancel, return_exceptions=True))
  244. for task in to_cancel:
  245. if task.cancelled():
  246. continue
  247. if task.exception() is not None:
  248. loop.call_exception_handler(
  249. {
  250. "message": "unhandled exception during asyncio.run() shutdown",
  251. "exception": task.exception(),
  252. "task": task,
  253. }
  254. )
  255. async def _shutdown_default_executor(loop: AbstractEventLoop) -> None:
  256. """Schedule the shutdown of the default executor."""
  257. def _do_shutdown(future: asyncio.futures.Future) -> None:
  258. try:
  259. loop._default_executor.shutdown(wait=True) # type: ignore[attr-defined]
  260. loop.call_soon_threadsafe(future.set_result, None)
  261. except Exception as ex:
  262. loop.call_soon_threadsafe(future.set_exception, ex)
  263. loop._executor_shutdown_called = True
  264. if loop._default_executor is None:
  265. return
  266. future = loop.create_future()
  267. thread = threading.Thread(target=_do_shutdown, args=(future,))
  268. thread.start()
  269. try:
  270. await future
  271. finally:
  272. thread.join()
  273. T_Retval = TypeVar("T_Retval")
  274. T_co = TypeVar("T_co", covariant=True)
  275. T_contra = TypeVar("T_contra", contravariant=True)
  276. PosArgsT = TypeVarTuple("PosArgsT")
  277. P = ParamSpec("P")
  278. _root_task: RunVar[asyncio.Task[Any] | None] = RunVar("_root_task")
  279. def find_root_task() -> asyncio.Task:
  280. root_task = _root_task.get(None)
  281. if root_task is not None and not root_task.done():
  282. return root_task
  283. # Look for a task that has been started via run_until_complete()
  284. for task in all_tasks():
  285. if task._callbacks and not task.done():
  286. for cb, context in task._callbacks:
  287. if (
  288. cb is _run_until_complete_cb
  289. or getattr(cb, "__module__", None) == "uvloop.loop"
  290. ):
  291. _root_task.set(task)
  292. def _unset(t: asyncio.Task[Any]) -> None:
  293. if vars := _run_vars.get(t.get_loop()):
  294. vars.pop(_root_task, None)
  295. # Register a callback to break the task -> loop -> _run_var[loop][_root_task] -> task cycle
  296. # Also run it in its own context to not create another reference.
  297. # We can't use RunVar.reset() here since these are called synchronously
  298. # and thus lowlevel.current_token() (which RunVar.reset() depends on) fails.
  299. task.add_done_callback(_unset, context=context)
  300. return task
  301. # Look up the topmost task in the AnyIO task tree, if possible
  302. task = cast(asyncio.Task, current_task())
  303. state = _task_states.get(task)
  304. if state:
  305. cancel_scope = state.cancel_scope
  306. while cancel_scope and cancel_scope._parent_scope is not None:
  307. cancel_scope = cancel_scope._parent_scope
  308. if cancel_scope is not None:
  309. return cast(asyncio.Task, cancel_scope._host_task)
  310. return task
  311. def _task_started(task: asyncio.Task) -> bool:
  312. """Return ``True`` if the task has been started and has not finished."""
  313. # The task coro should never be None here, as we never add finished tasks to the
  314. # task list
  315. coro = task.get_coro()
  316. assert coro is not None
  317. return getcoroutinestate(coro) in (CORO_RUNNING, CORO_SUSPENDED)
  318. #
  319. # Timeouts and cancellation
  320. #
  321. def is_anyio_cancellation(exc: CancelledError) -> bool:
  322. # Sometimes third party frameworks catch a CancelledError and raise a new one, so as
  323. # a workaround we have to look at the previous ones in __context__ too for a
  324. # matching cancel message
  325. while True:
  326. if (
  327. exc.args
  328. and isinstance(exc.args[0], str)
  329. and exc.args[0].startswith("Cancelled via cancel scope ")
  330. ):
  331. return True
  332. if isinstance(exc.__context__, CancelledError):
  333. exc = exc.__context__
  334. continue
  335. return False
  336. class CancelScope(BaseCancelScope):
  337. __slots__ = (
  338. "_active",
  339. "_cancel_called",
  340. "_cancel_handle",
  341. "_cancel_reason",
  342. "_cancelled_caught",
  343. "_child_scopes",
  344. "_deadline",
  345. "_host_task",
  346. "_parent_scope",
  347. "_pending_uncancellations",
  348. "_shield",
  349. "_tasks",
  350. "_timeout_handle",
  351. )
  352. def __new__(cls, *, deadline: float = math.inf, shield: bool = False) -> Self:
  353. return object.__new__(cls)
  354. def __init__(self, deadline: float = math.inf, shield: bool = False):
  355. self._deadline = deadline
  356. self._shield = shield
  357. self._parent_scope: CancelScope | None = None
  358. self._child_scopes: set[CancelScope] = set()
  359. self._cancel_called = False
  360. self._cancel_reason: str | None = None
  361. self._cancelled_caught = False
  362. self._active = False
  363. self._timeout_handle: asyncio.TimerHandle | None = None
  364. self._cancel_handle: asyncio.Handle | None = None
  365. self._tasks: set[asyncio.Task] = set()
  366. self._host_task: asyncio.Task | None = None
  367. if sys.version_info >= (3, 11):
  368. self._pending_uncancellations: int | None = 0
  369. else:
  370. self._pending_uncancellations = None
  371. def __enter__(self) -> Self:
  372. if self._active:
  373. raise RuntimeError(
  374. "Each CancelScope may only be used for a single 'with' block"
  375. )
  376. self._host_task = host_task = cast(asyncio.Task, current_task())
  377. self._tasks.add(host_task)
  378. try:
  379. task_state = _task_states[host_task]
  380. except KeyError:
  381. task_state = TaskState(None, self)
  382. _task_states[host_task] = task_state
  383. else:
  384. self._parent_scope = task_state.cancel_scope
  385. task_state.cancel_scope = self
  386. if self._parent_scope is not None:
  387. # If using an eager task factory, the parent scope may not even contain
  388. # the host task
  389. self._parent_scope._child_scopes.add(self)
  390. self._parent_scope._tasks.discard(host_task)
  391. self._timeout()
  392. self._active = True
  393. # Start cancelling the host task if the scope was cancelled before entering
  394. if self._cancel_called:
  395. self._deliver_cancellation(self)
  396. return self
  397. def __exit__(
  398. self,
  399. exc_type: type[BaseException] | None,
  400. exc_val: BaseException | None,
  401. exc_tb: TracebackType | None,
  402. ) -> bool:
  403. del exc_tb
  404. if not self._active:
  405. raise RuntimeError("This cancel scope is not active")
  406. if current_task() is not self._host_task:
  407. raise RuntimeError(
  408. "Attempted to exit cancel scope in a different task than it was "
  409. "entered in"
  410. )
  411. assert self._host_task is not None
  412. host_task_state = _task_states.get(self._host_task)
  413. if host_task_state is None or host_task_state.cancel_scope is not self:
  414. raise RuntimeError(
  415. "Attempted to exit a cancel scope that isn't the current tasks's "
  416. "current cancel scope"
  417. )
  418. try:
  419. self._active = False
  420. if self._timeout_handle:
  421. self._timeout_handle.cancel()
  422. self._timeout_handle = None
  423. self._tasks.remove(self._host_task)
  424. if self._parent_scope is not None:
  425. self._parent_scope._child_scopes.remove(self)
  426. self._parent_scope._tasks.add(self._host_task)
  427. host_task_state.cancel_scope = self._parent_scope
  428. # Restart the cancellation effort in the closest visible, cancelled parent
  429. # scope if necessary
  430. self._restart_cancellation_in_parent()
  431. # We only swallow the exception iff it was an AnyIO CancelledError, either
  432. # directly as exc_val or inside an exception group and there are no cancelled
  433. # parent cancel scopes visible to us here
  434. if self._cancel_called and not self._parent_cancellation_is_visible_to_us:
  435. # For each level-cancel() call made on the host task, call uncancel()
  436. while self._pending_uncancellations:
  437. self._host_task.uncancel()
  438. self._pending_uncancellations -= 1
  439. # Update cancelled_caught and check for exceptions we must not swallow
  440. if isinstance(exc_val, BaseExceptionGroup):
  441. cancelleds_caught, remaining = exc_val.split(
  442. lambda exc: (
  443. isinstance(exc, CancelledError)
  444. and is_anyio_cancellation(exc)
  445. )
  446. )
  447. if cancelleds_caught is None:
  448. return False
  449. self._cancelled_caught = True
  450. if remaining is None:
  451. return True
  452. context = remaining.__context__
  453. try:
  454. # Preserve __cause__ and __suppress_context__ by avoiding `raise
  455. # ... from ...`
  456. raise remaining
  457. finally:
  458. # Preserve __context__
  459. remaining.__context__ = context
  460. del context
  461. else:
  462. if isinstance(exc_val, CancelledError) and is_anyio_cancellation(
  463. exc_val
  464. ):
  465. self._cancelled_caught = True
  466. return True
  467. else:
  468. return False
  469. else:
  470. if self._pending_uncancellations:
  471. assert self._parent_scope is not None
  472. assert self._parent_scope._pending_uncancellations is not None
  473. self._parent_scope._pending_uncancellations += (
  474. self._pending_uncancellations
  475. )
  476. self._pending_uncancellations = 0
  477. return False
  478. finally:
  479. self._host_task = None
  480. del exc_val
  481. @property
  482. def _effectively_cancelled(self) -> bool:
  483. cancel_scope: CancelScope | None = self
  484. while cancel_scope is not None:
  485. if cancel_scope._cancel_called:
  486. return True
  487. if cancel_scope.shield:
  488. return False
  489. cancel_scope = cancel_scope._parent_scope
  490. return False
  491. @property
  492. def _parent_cancellation_is_visible_to_us(self) -> bool:
  493. return (
  494. self._parent_scope is not None
  495. and not self.shield
  496. and self._parent_scope._effectively_cancelled
  497. )
  498. def _timeout(self) -> None:
  499. if self._deadline != math.inf:
  500. loop = get_running_loop()
  501. if loop.time() >= self._deadline:
  502. self.cancel("deadline exceeded")
  503. else:
  504. self._timeout_handle = loop.call_at(self._deadline, self._timeout)
  505. def _deliver_cancellation(self, origin: CancelScope) -> bool:
  506. """
  507. Deliver cancellation to directly contained tasks and nested cancel scopes.
  508. Schedule another run at the end if we still have tasks eligible for
  509. cancellation.
  510. :param origin: the cancel scope that originated the cancellation
  511. :return: ``True`` if the delivery needs to be retried on the next cycle
  512. """
  513. should_retry = False
  514. current = current_task()
  515. for task in self._tasks:
  516. # Always skip tasks that are already done (see issue #1111)
  517. if task.done():
  518. continue
  519. should_retry = True
  520. if task._must_cancel: # type: ignore[attr-defined]
  521. continue
  522. # The task is eligible for cancellation if it has started
  523. if task is not current and (task is self._host_task or _task_started(task)):
  524. waiter = task._fut_waiter # type: ignore[attr-defined]
  525. if not isinstance(waiter, asyncio.Future) or not waiter.done():
  526. task.cancel(origin._cancel_reason)
  527. if (
  528. task is origin._host_task
  529. and origin._pending_uncancellations is not None
  530. ):
  531. origin._pending_uncancellations += 1
  532. # Deliver cancellation to child scopes that aren't shielded or running their own
  533. # cancellation callbacks
  534. for scope in self._child_scopes:
  535. if not scope._shield and not scope.cancel_called:
  536. should_retry = scope._deliver_cancellation(origin) or should_retry
  537. # Schedule another callback if there are still tasks left
  538. if origin is self:
  539. if should_retry:
  540. self._cancel_handle = get_running_loop().call_soon(
  541. self._deliver_cancellation, origin
  542. )
  543. else:
  544. self._cancel_handle = None
  545. return should_retry
  546. def _restart_cancellation_in_parent(self) -> None:
  547. """
  548. Restart the cancellation effort in the closest directly cancelled parent scope.
  549. """
  550. scope = self._parent_scope
  551. while scope is not None:
  552. if scope._cancel_called:
  553. if scope._cancel_handle is None:
  554. scope._deliver_cancellation(scope)
  555. break
  556. # No point in looking beyond any shielded scope
  557. if scope._shield:
  558. break
  559. scope = scope._parent_scope
  560. def _reparent(self, new_parent: CancelScope) -> None:
  561. """
  562. Move this active scope from its current parent to ``new_parent``.
  563. Used by :meth:`TaskGroup.start` to move a task that has just called
  564. ``task_status.started()`` into the target task group's cancel scope.
  565. """
  566. if self._parent_scope is new_parent:
  567. return
  568. if self._parent_scope is not None:
  569. self._parent_scope._child_scopes.discard(self)
  570. self._parent_scope = new_parent
  571. new_parent._child_scopes.add(self)
  572. # If the new parent (or an ancestor) is already cancelled, (re)start the
  573. # delivery loop to ensure we're cancelled at next checkpoint like Trio.
  574. self._restart_cancellation_in_parent()
  575. def cancel(self, reason: str | None = None) -> None:
  576. if not self._cancel_called:
  577. if self._timeout_handle:
  578. self._timeout_handle.cancel()
  579. self._timeout_handle = None
  580. self._cancel_called = True
  581. self._cancel_reason = f"Cancelled via cancel scope {id(self):x}"
  582. if task := current_task():
  583. self._cancel_reason += f" by {task}"
  584. if reason:
  585. self._cancel_reason += f"; reason: {reason}"
  586. if self._host_task is not None:
  587. self._deliver_cancellation(self)
  588. @property
  589. def deadline(self) -> float:
  590. return self._deadline
  591. @deadline.setter
  592. def deadline(self, value: float) -> None:
  593. self._deadline = float(value)
  594. if self._timeout_handle is not None:
  595. self._timeout_handle.cancel()
  596. self._timeout_handle = None
  597. if self._active and not self._cancel_called:
  598. self._timeout()
  599. @property
  600. def cancel_called(self) -> bool:
  601. return self._cancel_called
  602. @property
  603. def cancelled_caught(self) -> bool:
  604. return self._cancelled_caught
  605. @property
  606. def shield(self) -> bool:
  607. return self._shield
  608. @shield.setter
  609. def shield(self, value: bool) -> None:
  610. if self._shield != value:
  611. self._shield = value
  612. if not value:
  613. self._restart_cancellation_in_parent()
  614. #
  615. # Task states
  616. #
  617. class TaskState:
  618. """
  619. Encapsulates auxiliary task information that cannot be added to the Task instance
  620. itself because there are no guarantees about its implementation.
  621. """
  622. __slots__ = "__weakref__", "cancel_scope", "parent_id"
  623. def __init__(self, parent_id: int | None, cancel_scope: CancelScope | None):
  624. self.parent_id = parent_id
  625. self.cancel_scope = cancel_scope
  626. _task_states: WeakKeyDictionary[asyncio.Task, TaskState] = WeakKeyDictionary()
  627. #
  628. # Task groups
  629. #
  630. class _AsyncioTaskStatus(abc.TaskStatus[T_contra]):
  631. def __init__(
  632. self,
  633. future: asyncio.Future,
  634. parent_id: int,
  635. target_scope: CancelScope,
  636. spawn_scope: CancelScope,
  637. ):
  638. self._future = future
  639. self._parent_id = parent_id
  640. # The eventual parent scope for this spawn_scope
  641. # (after task_status.started() has been called)
  642. self._target_scope = target_scope
  643. # The task's own cancel scope, also held by its TaskHandle
  644. self._spawn_scope = spawn_scope
  645. def started(self, value: T_contra | None = None) -> None:
  646. task = cast(asyncio.Task, current_task())
  647. _task_states[task].parent_id = self._parent_id
  648. if self._future.done():
  649. if not self._future.cancelled():
  650. raise RuntimeError("called 'started' twice on the same task status")
  651. else:
  652. # Caller of start() was cancelled, nothing to reparent
  653. return
  654. self._future.set_result(value)
  655. self._spawn_scope._reparent(self._target_scope)
  656. if sys.version_info >= (3, 12):
  657. _eager_task_factory_code: CodeType | None = asyncio.eager_task_factory.__code__
  658. else:
  659. _eager_task_factory_code = None
  660. class TaskGroup(abc.TaskGroup):
  661. def __init__(self) -> None:
  662. self.cancel_scope: CancelScope = CancelScope()
  663. self._entered = False
  664. self._exceptions: list[BaseException] = []
  665. self._tasks: set[asyncio.Task] = set()
  666. self._on_completed_fut: asyncio.Future[None] | None = None
  667. async def __aenter__(self) -> Self:
  668. if self._entered:
  669. raise RuntimeError("TaskGroup cannot be entered more than once")
  670. self._entered = True
  671. self.cancel_scope.__enter__()
  672. return self
  673. async def __aexit__(
  674. self,
  675. exc_type: type[BaseException] | None,
  676. exc_val: BaseException | None,
  677. exc_tb: TracebackType | None,
  678. ) -> bool:
  679. try:
  680. if exc_val is not None:
  681. self.cancel_scope.cancel()
  682. if not isinstance(exc_val, CancelledError):
  683. self._exceptions.append(exc_val)
  684. loop = get_running_loop()
  685. try:
  686. if self._tasks:
  687. with CancelScope() as wait_scope:
  688. while self._tasks:
  689. self._on_completed_fut = loop.create_future()
  690. try:
  691. await self._on_completed_fut
  692. except CancelledError as exc:
  693. # Shield the scope against further cancellation attempts,
  694. # as they're not productive (#695)
  695. wait_scope.shield = True
  696. self.cancel_scope.cancel()
  697. # Set exc_val from the cancellation exception if it was
  698. # previously unset. However, we should not replace a native
  699. # cancellation exception with one raise by a cancel scope.
  700. if exc_val is None or (
  701. isinstance(exc_val, CancelledError)
  702. and not is_anyio_cancellation(exc)
  703. ):
  704. exc_val = exc
  705. self._on_completed_fut = None
  706. else:
  707. # If there are no child tasks to wait on, run at least one checkpoint
  708. # anyway
  709. await AsyncIOBackend.cancel_shielded_checkpoint()
  710. if self._exceptions:
  711. # The exception that got us here should already have been
  712. # added to self._exceptions so it's ok to break exception
  713. # chaining and avoid adding a "During handling of above..."
  714. # for each nesting level.
  715. raise BaseExceptionGroup(
  716. "unhandled errors in a TaskGroup", self._exceptions
  717. ) from None
  718. elif exc_val:
  719. raise exc_val
  720. except BaseException as exc:
  721. if self.cancel_scope.__exit__(type(exc), exc, exc.__traceback__):
  722. return True
  723. raise
  724. return self.cancel_scope.__exit__(exc_type, exc_val, exc_tb)
  725. finally:
  726. del exc_val, exc_tb, self._exceptions
  727. def _spawn(
  728. self,
  729. coro: Coroutine[Any, Any, T_co],
  730. name: object,
  731. task_status: _AsyncioTaskStatus | None = None,
  732. ) -> TaskHandle[T_co]:
  733. task_status_future = task_status._future if task_status is not None else None
  734. def task_done(_task: asyncio.Task) -> None:
  735. if sys.version_info >= (3, 14) and self.cancel_scope._host_task is not None:
  736. asyncio.future_discard_from_awaited_by(
  737. _task, self.cancel_scope._host_task
  738. )
  739. task_state = _task_states[_task]
  740. assert task_state.cancel_scope is not None
  741. assert _task in task_state.cancel_scope._tasks
  742. task_state.cancel_scope._tasks.remove(_task)
  743. self._tasks.remove(task)
  744. del _task_states[_task]
  745. if self._on_completed_fut is not None and not self._tasks:
  746. try:
  747. self._on_completed_fut.set_result(None)
  748. except asyncio.InvalidStateError:
  749. pass
  750. try:
  751. exc = _task.exception()
  752. except CancelledError as e:
  753. while isinstance(e.__context__, CancelledError):
  754. e = e.__context__
  755. exc = e
  756. if exc is not None:
  757. # The future can only be in the cancelled state if the host task was
  758. # cancelled, so return immediately instead of adding one more
  759. # CancelledError to the exceptions list
  760. if task_status_future is not None and task_status_future.cancelled():
  761. return
  762. if task_status_future is None or task_status_future.done():
  763. if not isinstance(exc, CancelledError):
  764. self._exceptions.append(exc)
  765. if not self.cancel_scope._effectively_cancelled:
  766. self.cancel_scope.cancel()
  767. else:
  768. task_status_future.set_exception(exc)
  769. elif task_status_future is not None and not task_status_future.done():
  770. task_status_future.set_exception(
  771. RuntimeError("Child exited without calling task_status.started()")
  772. )
  773. if task_status_future is not None:
  774. parent_id = id(current_task())
  775. caller_state = _task_states.get(cast(asyncio.Task, current_task()))
  776. if caller_state is not None and caller_state.cancel_scope is not None:
  777. initial_scope = caller_state.cancel_scope
  778. else:
  779. # The caller is an unmanaged task (no task state)
  780. initial_scope = self.cancel_scope
  781. else:
  782. parent_id = id(self.cancel_scope._host_task)
  783. initial_scope = self.cancel_scope
  784. spawn_scope = task_status._spawn_scope if task_status is not None else None
  785. handle = TaskHandle(coro, name, cancel_scope=spawn_scope)
  786. loop = asyncio.get_running_loop()
  787. wrapper_coro = handle._run_coro()
  788. try:
  789. if (
  790. (factory := loop.get_task_factory())
  791. and getattr(factory, "__code__", None) is _eager_task_factory_code
  792. and (closure := getattr(factory, "__closure__", None))
  793. ):
  794. custom_task_constructor = closure[0].cell_contents
  795. task = custom_task_constructor(
  796. wrapper_coro, loop=loop, name=handle.name
  797. )
  798. else:
  799. task = loop.create_task(wrapper_coro, name=handle.name)
  800. except BaseException:
  801. with suppress(BaseException):
  802. wrapper_coro.close()
  803. with suppress(BaseException):
  804. coro.close()
  805. raise
  806. # Make the spawned task inherit the initial cancel scope
  807. _task_states[task] = TaskState(parent_id=parent_id, cancel_scope=initial_scope)
  808. initial_scope._tasks.add(task)
  809. self._tasks.add(task)
  810. if sys.version_info >= (3, 14) and self.cancel_scope._host_task is not None:
  811. asyncio.future_add_to_awaited_by(task, self.cancel_scope._host_task)
  812. task.add_done_callback(task_done)
  813. return handle
  814. def create_task(
  815. self,
  816. coro: Coroutine[Any, Any, T_co],
  817. *,
  818. name: object = None,
  819. context: Context | None = None,
  820. ) -> TaskHandle[T_co]:
  821. if not isinstance(coro, Coroutine):
  822. raise TypeError(f"expected a coroutine, got {coro.__class__.__qualname__}")
  823. if not self._entered or not self.cancel_scope._active:
  824. coro.close()
  825. raise RuntimeError(
  826. "This task group is not active; no new tasks can be started."
  827. )
  828. final_name = get_coro_name(coro, name)
  829. if context is not None:
  830. return context.run(self._spawn, coro, name=final_name)
  831. else:
  832. return self._spawn(coro, name=final_name)
  833. async def start(
  834. self,
  835. func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
  836. *args: Unpack[PosArgsT],
  837. name: object = None,
  838. return_handle: Literal[False, True] = False,
  839. ) -> Any:
  840. if not self._entered or not self.cancel_scope._active:
  841. raise RuntimeError(
  842. "This task group is not active; no new tasks can be started."
  843. )
  844. # Until task_status.started() is called, a task spawned via start() belongs to
  845. # the *caller's* cancel scope, not the target group's, so cancelling the group
  846. # does not cancel a task that hasn't reported startup yet. The
  847. # task_status.started() call moves the task's cancel scope to this task group's
  848. # cancel scope.
  849. #
  850. # The caller may be an unmanaged task (no _task_states entry), in which case
  851. # fall back to the group's own scope.
  852. future: asyncio.Future = asyncio.Future()
  853. final_name = get_callable_name(func, name)
  854. task_status: _AsyncioTaskStatus[Any] = _AsyncioTaskStatus(
  855. future, id(self.cancel_scope._host_task), self.cancel_scope, CancelScope()
  856. )
  857. coro = call_for_coroutine(func, args, task_status=task_status)
  858. handle = self._spawn(coro, final_name, task_status)
  859. # If the task raises an exception after sending a start value without a switch
  860. # point between, the task group is cancelled and this method never proceeds to
  861. # process the completed future. That's why we have to have a shielded cancel
  862. # scope here.
  863. try:
  864. await future
  865. except BaseException:
  866. if handle.status is TaskHandle.Status.PENDING:
  867. # Cancel the task and wait for it to exit before returning
  868. handle.cancel()
  869. with CancelScope(shield=True):
  870. await handle.wait()
  871. raise
  872. if return_handle:
  873. handle._start_value = future.result()
  874. return handle
  875. else:
  876. return future.result()
  877. #
  878. # Threads
  879. #
  880. _Retval_Queue_Type = tuple[T_Retval | None, BaseException | None]
  881. class WorkerThread(Thread):
  882. MAX_IDLE_TIME = 10 # seconds
  883. def __init__(
  884. self,
  885. root_task: asyncio.Task,
  886. workers: set[WorkerThread],
  887. idle_workers: deque[WorkerThread],
  888. ):
  889. kwargs: dict[str, Any] = {}
  890. if sys.version_info >= (3, 14):
  891. kwargs["context"] = Context()
  892. super().__init__(name="AnyIO worker thread", **kwargs)
  893. self.root_task = root_task
  894. self.workers = workers
  895. self.idle_workers = idle_workers
  896. self.loop = root_task._loop
  897. self.queue: Queue[
  898. tuple[Context, Callable, tuple, asyncio.Future, CancelScope] | None
  899. ] = Queue(2)
  900. self.idle_since = AsyncIOBackend.current_time()
  901. self.stopping = False
  902. def _report_result(
  903. self, future: asyncio.Future, result: Any, exc: BaseException | None
  904. ) -> None:
  905. self.idle_since = AsyncIOBackend.current_time()
  906. if not self.stopping:
  907. self.idle_workers.append(self)
  908. if not future.cancelled():
  909. if exc is not None:
  910. if isinstance(exc, StopIteration):
  911. new_exc = RuntimeError("coroutine raised StopIteration")
  912. new_exc.__cause__ = exc
  913. exc = new_exc
  914. future.set_exception(exc)
  915. else:
  916. future.set_result(result)
  917. def run(self) -> None:
  918. with claim_worker_thread(AsyncIOBackend, self.loop):
  919. while True:
  920. item = self.queue.get()
  921. if item is None:
  922. # Shutdown command received
  923. return
  924. context, func, args, future, cancel_scope = item
  925. if not future.cancelled():
  926. result = None
  927. exception: BaseException | None = None
  928. threadlocals.current_cancel_scope = cancel_scope
  929. try:
  930. result = context.run(func, *args)
  931. except BaseException as exc:
  932. exception = exc
  933. finally:
  934. del threadlocals.current_cancel_scope
  935. try:
  936. self.loop.call_soon_threadsafe(
  937. self._report_result, future, result, exception
  938. )
  939. except RuntimeError:
  940. if not self.loop.is_closed():
  941. raise
  942. del result, exception
  943. self.queue.task_done()
  944. del item, context, func, args, future, cancel_scope
  945. def stop(self, f: asyncio.Task | None = None) -> None:
  946. self.stopping = True
  947. self.queue.put_nowait(None)
  948. self.workers.discard(self)
  949. try:
  950. self.idle_workers.remove(self)
  951. except ValueError:
  952. pass
  953. _threadpool_idle_workers: RunVar[deque[WorkerThread]] = RunVar(
  954. "_threadpool_idle_workers"
  955. )
  956. _threadpool_workers: RunVar[set[WorkerThread]] = RunVar("_threadpool_workers")
  957. #
  958. # Subprocesses
  959. #
  960. @dataclass(eq=False)
  961. class StreamReaderWrapper(abc.ByteReceiveStream):
  962. _stream: asyncio.StreamReader
  963. async def receive(self, max_bytes: int = 65536) -> bytes:
  964. if max_bytes < 1:
  965. raise ValueError("max_bytes must be a positive integer")
  966. data = await self._stream.read(max_bytes)
  967. if data:
  968. return data
  969. else:
  970. raise EndOfStream
  971. async def aclose(self) -> None:
  972. self._stream.set_exception(ClosedResourceError())
  973. await AsyncIOBackend.checkpoint()
  974. @dataclass(eq=False)
  975. class StreamWriterWrapper(abc.ByteSendStream):
  976. _stream: asyncio.StreamWriter
  977. _closed: bool = field(init=False, default=False)
  978. async def send(self, item: bytes) -> None:
  979. await AsyncIOBackend.checkpoint_if_cancelled()
  980. stream_paused = self._stream._protocol._paused # type: ignore[attr-defined]
  981. try:
  982. self._stream.write(item)
  983. await self._stream.drain()
  984. except (ConnectionResetError, BrokenPipeError, RuntimeError) as exc:
  985. # If closed by us and/or the peer:
  986. # * on stdlib, drain() raises ConnectionResetError or BrokenPipeError
  987. # * on uvloop and Winloop, write() eventually starts raising RuntimeError
  988. if self._closed:
  989. raise ClosedResourceError from exc
  990. elif self._stream.is_closing():
  991. raise BrokenResourceError from exc
  992. raise
  993. if not stream_paused:
  994. await AsyncIOBackend.cancel_shielded_checkpoint()
  995. async def aclose(self) -> None:
  996. self._closed = True
  997. self._stream.close()
  998. await AsyncIOBackend.checkpoint()
  999. @dataclass(eq=False)
  1000. class Process(abc.Process):
  1001. _process: asyncio.subprocess.Process
  1002. _stdin: StreamWriterWrapper | None
  1003. _stdout: StreamReaderWrapper | None
  1004. _stderr: StreamReaderWrapper | None
  1005. _exited: asyncio.Event
  1006. _transport: asyncio.SubprocessTransport
  1007. async def aclose(self) -> None:
  1008. with CancelScope(shield=True) as scope:
  1009. # We need to close the underlying pipe_transports as well to allow a
  1010. # process blocking on full buffers to receive SIGPIPE and exit.
  1011. if self._stdin:
  1012. await self._stdin.aclose()
  1013. if pipe := self._transport.get_pipe_transport(0):
  1014. pipe.close()
  1015. if self._stdout:
  1016. await self._stdout.aclose()
  1017. if pipe := self._transport.get_pipe_transport(1):
  1018. pipe.close()
  1019. if self._stderr:
  1020. await self._stderr.aclose()
  1021. if pipe := self._transport.get_pipe_transport(2):
  1022. pipe.close()
  1023. scope.shield = False
  1024. try:
  1025. await self.wait()
  1026. except BaseException:
  1027. scope.shield = True
  1028. # Closing the transport on asyncio also handles sending kill
  1029. self._transport.close()
  1030. await self.wait()
  1031. raise
  1032. async def wait(self) -> int:
  1033. await self._exited.wait()
  1034. assert self._process.returncode is not None
  1035. return self._process.returncode
  1036. def terminate(self) -> None:
  1037. self._process.terminate()
  1038. def kill(self) -> None:
  1039. self._process.kill()
  1040. def send_signal(self, signal: int) -> None:
  1041. self._process.send_signal(signal)
  1042. @property
  1043. def pid(self) -> int:
  1044. return self._process.pid
  1045. @property
  1046. def returncode(self) -> int | None:
  1047. return self._process.returncode
  1048. @property
  1049. def stdin(self) -> abc.ByteSendStream | None:
  1050. return self._stdin
  1051. @property
  1052. def stdout(self) -> abc.ByteReceiveStream | None:
  1053. return self._stdout
  1054. @property
  1055. def stderr(self) -> abc.ByteReceiveStream | None:
  1056. return self._stderr
  1057. def _forcibly_shutdown_process_pool_on_exit(
  1058. workers: set[Process], _task: object
  1059. ) -> None:
  1060. """
  1061. Forcibly shuts down worker processes belonging to this event loop."""
  1062. child_watcher: asyncio.AbstractChildWatcher | None = None # type: ignore[name-defined]
  1063. if sys.version_info < (3, 12):
  1064. try:
  1065. child_watcher = asyncio.get_event_loop_policy().get_child_watcher()
  1066. except NotImplementedError:
  1067. pass
  1068. # Close as much as possible (w/o async/await) to avoid warnings
  1069. for process in workers.copy():
  1070. if process.returncode is not None:
  1071. continue
  1072. process._stdin._stream._transport.close() # type: ignore[union-attr]
  1073. process._stdout._stream._transport.close() # type: ignore[union-attr]
  1074. process._stderr._stream._transport.close() # type: ignore[union-attr]
  1075. process.kill()
  1076. if child_watcher:
  1077. child_watcher.remove_child_handler(process.pid)
  1078. async def _shutdown_process_pool_on_exit(workers: set[abc.Process]) -> None:
  1079. """
  1080. Shuts down worker processes belonging to this event loop.
  1081. NOTE: this only works when the event loop was started using asyncio.run() or
  1082. anyio.run().
  1083. """
  1084. process: abc.Process
  1085. try:
  1086. await sleep(math.inf)
  1087. except asyncio.CancelledError:
  1088. workers = workers.copy()
  1089. for process in workers:
  1090. if process.returncode is None:
  1091. process.kill()
  1092. for process in workers:
  1093. await process.aclose()
  1094. #
  1095. # Sockets and networking
  1096. #
  1097. class StreamProtocol(asyncio.Protocol):
  1098. read_queue: deque[bytes]
  1099. read_event: asyncio.Event
  1100. write_event: asyncio.Event
  1101. exception: Exception | None = None
  1102. is_at_eof: bool = False
  1103. def connection_made(self, transport: asyncio.BaseTransport) -> None:
  1104. self.read_queue = deque()
  1105. self.read_event = asyncio.Event()
  1106. self.write_event = asyncio.Event()
  1107. self.write_event.set()
  1108. cast(asyncio.Transport, transport).set_write_buffer_limits(0)
  1109. def connection_lost(self, exc: Exception | None) -> None:
  1110. if exc:
  1111. self.exception = exc
  1112. self.read_event.set()
  1113. self.write_event.set()
  1114. def data_received(self, data: bytes) -> None:
  1115. # ProactorEventloop sometimes sends bytearray instead of bytes
  1116. self.read_queue.append(bytes(data))
  1117. self.read_event.set()
  1118. def eof_received(self) -> bool | None:
  1119. self.is_at_eof = True
  1120. self.read_event.set()
  1121. return True
  1122. def pause_writing(self) -> None:
  1123. self.write_event = asyncio.Event()
  1124. def resume_writing(self) -> None:
  1125. self.write_event.set()
  1126. class DatagramProtocol(asyncio.DatagramProtocol):
  1127. read_queue: deque[tuple[bytes, IPSockAddrType]]
  1128. read_event: asyncio.Event
  1129. write_event: asyncio.Event
  1130. closed_event: asyncio.Event
  1131. exception: Exception | None = None
  1132. def connection_made(self, transport: asyncio.BaseTransport) -> None:
  1133. self.read_queue = deque(maxlen=100) # arbitrary value
  1134. self.read_event = asyncio.Event()
  1135. self.write_event = asyncio.Event()
  1136. self.closed_event = asyncio.Event()
  1137. self.write_event.set()
  1138. def connection_lost(self, exc: Exception | None) -> None:
  1139. self.read_event.set()
  1140. self.write_event.set()
  1141. self.closed_event.set()
  1142. def datagram_received(self, data: bytes, addr: IPSockAddrType) -> None:
  1143. addr = convert_ipv6_sockaddr(addr)
  1144. self.read_queue.append((data, addr))
  1145. self.read_event.set()
  1146. def error_received(self, exc: Exception) -> None:
  1147. self.exception = exc
  1148. def pause_writing(self) -> None:
  1149. self.write_event.clear()
  1150. def resume_writing(self) -> None:
  1151. self.write_event.set()
  1152. class SocketStream(abc.SocketStream):
  1153. def __init__(self, transport: asyncio.Transport, protocol: StreamProtocol):
  1154. self._transport = transport
  1155. self._protocol = protocol
  1156. self._receive_guard = ResourceGuard("reading from")
  1157. self._send_guard = ResourceGuard("writing to")
  1158. self._closed = False
  1159. @property
  1160. def _raw_socket(self) -> socket.socket:
  1161. return self._transport.get_extra_info("socket")
  1162. async def receive(self, max_bytes: int = 65536) -> bytes:
  1163. if max_bytes < 1:
  1164. raise ValueError("max_bytes must be a positive integer")
  1165. with self._receive_guard:
  1166. if (
  1167. not self._protocol.read_event.is_set()
  1168. and not self._transport.is_closing()
  1169. and not self._protocol.is_at_eof
  1170. ):
  1171. self._transport.resume_reading()
  1172. await self._protocol.read_event.wait()
  1173. self._transport.pause_reading()
  1174. else:
  1175. await AsyncIOBackend.checkpoint()
  1176. try:
  1177. chunk = self._protocol.read_queue.popleft()
  1178. except IndexError:
  1179. if self._closed:
  1180. raise ClosedResourceError from None
  1181. elif self._protocol.exception:
  1182. raise BrokenResourceError from self._protocol.exception
  1183. else:
  1184. raise EndOfStream from None
  1185. if len(chunk) > max_bytes:
  1186. # Split the oversized chunk
  1187. chunk, leftover = chunk[:max_bytes], chunk[max_bytes:]
  1188. self._protocol.read_queue.appendleft(leftover)
  1189. # If the read queue is empty, clear the flag so that the next call will
  1190. # block until data is available
  1191. if not self._protocol.read_queue:
  1192. self._protocol.read_event.clear()
  1193. return chunk
  1194. async def send(self, item: bytes) -> None:
  1195. with self._send_guard:
  1196. await AsyncIOBackend.checkpoint()
  1197. if self._closed:
  1198. raise ClosedResourceError
  1199. elif self._protocol.exception is not None:
  1200. raise BrokenResourceError from self._protocol.exception
  1201. try:
  1202. self._transport.write(item)
  1203. except RuntimeError as exc:
  1204. if self._transport.is_closing():
  1205. raise BrokenResourceError from exc
  1206. else:
  1207. raise
  1208. await self._protocol.write_event.wait()
  1209. async def send_eof(self) -> None:
  1210. try:
  1211. self._transport.write_eof()
  1212. except OSError:
  1213. pass
  1214. async def aclose(self) -> None:
  1215. self._closed = True
  1216. if not self._transport.is_closing():
  1217. try:
  1218. self._transport.write_eof()
  1219. except OSError:
  1220. pass
  1221. self._transport.close()
  1222. await sleep(0)
  1223. self._transport.abort()
  1224. class _RawSocketMixin:
  1225. _receive_future: asyncio.Future | None = None
  1226. _send_future: asyncio.Future | None = None
  1227. _closing = False
  1228. def __init__(self, raw_socket: socket.socket):
  1229. self.__raw_socket = raw_socket
  1230. self._receive_guard = ResourceGuard("reading from")
  1231. self._send_guard = ResourceGuard("writing to")
  1232. @property
  1233. def _raw_socket(self) -> socket.socket:
  1234. return self.__raw_socket
  1235. def _wait_until_readable(self, loop: asyncio.AbstractEventLoop) -> asyncio.Future:
  1236. def callback(f: object) -> None:
  1237. del self._receive_future
  1238. loop.remove_reader(self.__raw_socket)
  1239. f = self._receive_future = asyncio.Future()
  1240. loop.add_reader(self.__raw_socket, f.set_result, None)
  1241. f.add_done_callback(callback)
  1242. return f
  1243. def _wait_until_writable(self, loop: asyncio.AbstractEventLoop) -> asyncio.Future:
  1244. def callback(f: object) -> None:
  1245. del self._send_future
  1246. loop.remove_writer(self.__raw_socket)
  1247. f = self._send_future = asyncio.Future()
  1248. loop.add_writer(self.__raw_socket, f.set_result, None)
  1249. f.add_done_callback(callback)
  1250. return f
  1251. async def aclose(self) -> None:
  1252. if not self._closing:
  1253. self._closing = True
  1254. if self.__raw_socket.fileno() != -1:
  1255. self.__raw_socket.close()
  1256. if self._receive_future and not self._receive_future.done():
  1257. self._receive_future.set_result(None)
  1258. if self._send_future and not self._send_future.done():
  1259. self._send_future.set_result(None)
  1260. class UNIXSocketStream(_RawSocketMixin, abc.UNIXSocketStream):
  1261. async def send_eof(self) -> None:
  1262. with self._send_guard:
  1263. self._raw_socket.shutdown(socket.SHUT_WR)
  1264. async def receive(self, max_bytes: int = 65536) -> bytes:
  1265. if max_bytes < 1:
  1266. raise ValueError("max_bytes must be a positive integer")
  1267. loop = get_running_loop()
  1268. await AsyncIOBackend.checkpoint()
  1269. with self._receive_guard:
  1270. while True:
  1271. try:
  1272. data = self._raw_socket.recv(max_bytes)
  1273. except BlockingIOError:
  1274. await self._wait_until_readable(loop)
  1275. except OSError as exc:
  1276. if self._closing:
  1277. raise ClosedResourceError from None
  1278. else:
  1279. raise BrokenResourceError from exc
  1280. else:
  1281. if not data:
  1282. raise EndOfStream
  1283. return data
  1284. async def send(self, item: bytes) -> None:
  1285. loop = get_running_loop()
  1286. await AsyncIOBackend.checkpoint()
  1287. with self._send_guard:
  1288. view = memoryview(item)
  1289. while view:
  1290. try:
  1291. bytes_sent = self._raw_socket.send(view)
  1292. except BlockingIOError:
  1293. await self._wait_until_writable(loop)
  1294. except OSError as exc:
  1295. if self._closing:
  1296. raise ClosedResourceError from None
  1297. else:
  1298. raise BrokenResourceError from exc
  1299. else:
  1300. view = view[bytes_sent:]
  1301. async def receive_fds(self, msglen: int, maxfds: int) -> tuple[bytes, list[int]]:
  1302. if not isinstance(msglen, int) or msglen < 0:
  1303. raise ValueError("msglen must be a non-negative integer")
  1304. if not isinstance(maxfds, int) or maxfds < 1:
  1305. raise ValueError("maxfds must be a positive integer")
  1306. loop = get_running_loop()
  1307. fds = array.array("i")
  1308. await AsyncIOBackend.checkpoint()
  1309. with self._receive_guard:
  1310. while True:
  1311. try:
  1312. message, ancdata, _flags, _addr = self._raw_socket.recvmsg(
  1313. msglen, socket.CMSG_LEN(maxfds * fds.itemsize)
  1314. )
  1315. except BlockingIOError:
  1316. await self._wait_until_readable(loop)
  1317. except OSError as exc:
  1318. if self._closing:
  1319. raise ClosedResourceError from None
  1320. else:
  1321. raise BrokenResourceError from exc
  1322. else:
  1323. if not message and not ancdata:
  1324. raise EndOfStream
  1325. break
  1326. for cmsg_level, cmsg_type, cmsg_data in ancdata:
  1327. if cmsg_level != socket.SOL_SOCKET or cmsg_type != socket.SCM_RIGHTS:
  1328. raise RuntimeError(
  1329. f"Received unexpected ancillary data; message = {message!r}, "
  1330. f"cmsg_level = {cmsg_level}, cmsg_type = {cmsg_type}"
  1331. )
  1332. fds.frombytes(cmsg_data[: len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
  1333. return message, list(fds)
  1334. async def send_fds(self, message: bytes, fds: Collection[int | IOBase]) -> None:
  1335. if not message:
  1336. raise ValueError("message must not be empty")
  1337. if not fds:
  1338. raise ValueError("fds must not be empty")
  1339. loop = get_running_loop()
  1340. filenos: list[int] = []
  1341. for fd in fds:
  1342. if isinstance(fd, int):
  1343. filenos.append(fd)
  1344. elif isinstance(fd, IOBase):
  1345. filenos.append(fd.fileno())
  1346. fdarray = array.array("i", filenos)
  1347. await AsyncIOBackend.checkpoint()
  1348. with self._send_guard:
  1349. while True:
  1350. try:
  1351. # The ignore can be removed after mypy picks up
  1352. # https://github.com/python/typeshed/pull/5545
  1353. self._raw_socket.sendmsg(
  1354. [message], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fdarray)]
  1355. )
  1356. break
  1357. except BlockingIOError:
  1358. await self._wait_until_writable(loop)
  1359. except OSError as exc:
  1360. if self._closing:
  1361. raise ClosedResourceError from None
  1362. else:
  1363. raise BrokenResourceError from exc
  1364. class TCPSocketListener(abc.SocketListener):
  1365. _accept_scope: CancelScope | None = None
  1366. _closed = False
  1367. def __init__(self, raw_socket: socket.socket):
  1368. self.__raw_socket = raw_socket
  1369. self._loop = cast(asyncio.BaseEventLoop, get_running_loop())
  1370. self._accept_guard = ResourceGuard("accepting connections from")
  1371. @property
  1372. def _raw_socket(self) -> socket.socket:
  1373. return self.__raw_socket
  1374. async def accept(self) -> abc.SocketStream:
  1375. if self._closed:
  1376. raise ClosedResourceError
  1377. with self._accept_guard:
  1378. await AsyncIOBackend.checkpoint()
  1379. with CancelScope() as self._accept_scope:
  1380. try:
  1381. client_sock, _addr = await self._loop.sock_accept(self._raw_socket)
  1382. except asyncio.CancelledError:
  1383. # Workaround for https://bugs.python.org/issue41317
  1384. try:
  1385. self._loop.remove_reader(self._raw_socket)
  1386. except (ValueError, NotImplementedError):
  1387. pass
  1388. if self._closed:
  1389. raise ClosedResourceError from None
  1390. raise
  1391. finally:
  1392. self._accept_scope = None
  1393. client_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  1394. transport, protocol = await self._loop.connect_accepted_socket(
  1395. StreamProtocol, client_sock
  1396. )
  1397. return SocketStream(transport, protocol)
  1398. async def aclose(self) -> None:
  1399. if self._closed:
  1400. return
  1401. self._closed = True
  1402. if self._accept_scope:
  1403. # Workaround for https://bugs.python.org/issue41317
  1404. try:
  1405. self._loop.remove_reader(self._raw_socket)
  1406. except (ValueError, NotImplementedError):
  1407. pass
  1408. self._accept_scope.cancel()
  1409. await sleep(0)
  1410. self._raw_socket.close()
  1411. class UNIXSocketListener(abc.SocketListener):
  1412. def __init__(self, raw_socket: socket.socket):
  1413. self.__raw_socket = raw_socket
  1414. self._loop = get_running_loop()
  1415. self._accept_guard = ResourceGuard("accepting connections from")
  1416. self._closed = False
  1417. async def accept(self) -> abc.SocketStream:
  1418. await AsyncIOBackend.checkpoint()
  1419. with self._accept_guard:
  1420. while True:
  1421. try:
  1422. client_sock, _ = self.__raw_socket.accept()
  1423. client_sock.setblocking(False)
  1424. return UNIXSocketStream(client_sock)
  1425. except BlockingIOError:
  1426. f: asyncio.Future = asyncio.Future()
  1427. self._loop.add_reader(self.__raw_socket, f.set_result, None)
  1428. f.add_done_callback(
  1429. lambda _: self._loop.remove_reader(self.__raw_socket)
  1430. )
  1431. await f
  1432. except OSError as exc:
  1433. if self._closed:
  1434. raise ClosedResourceError from None
  1435. else:
  1436. raise BrokenResourceError from exc
  1437. async def aclose(self) -> None:
  1438. self._closed = True
  1439. self.__raw_socket.close()
  1440. @property
  1441. def _raw_socket(self) -> socket.socket:
  1442. return self.__raw_socket
  1443. class UDPSocket(abc.UDPSocket):
  1444. def __init__(
  1445. self, transport: asyncio.DatagramTransport, protocol: DatagramProtocol
  1446. ):
  1447. self._transport = transport
  1448. self._protocol = protocol
  1449. self._receive_guard = ResourceGuard("reading from")
  1450. self._send_guard = ResourceGuard("writing to")
  1451. self._closed = False
  1452. @property
  1453. def _raw_socket(self) -> socket.socket:
  1454. return self._transport.get_extra_info("socket")
  1455. async def aclose(self) -> None:
  1456. self._closed = True
  1457. if not self._transport.is_closing():
  1458. self._transport.close()
  1459. await self._protocol.closed_event.wait()
  1460. async def receive(self) -> tuple[bytes, IPSockAddrType]:
  1461. with self._receive_guard:
  1462. await AsyncIOBackend.checkpoint()
  1463. # If the buffer is empty, ask for more data
  1464. if not self._protocol.read_queue and not self._transport.is_closing():
  1465. self._protocol.read_event.clear()
  1466. await self._protocol.read_event.wait()
  1467. try:
  1468. return self._protocol.read_queue.popleft()
  1469. except IndexError:
  1470. if self._closed:
  1471. raise ClosedResourceError from None
  1472. else:
  1473. raise BrokenResourceError from None
  1474. async def send(self, item: UDPPacketType) -> None:
  1475. with self._send_guard:
  1476. await AsyncIOBackend.checkpoint()
  1477. await self._protocol.write_event.wait()
  1478. if self._closed:
  1479. raise ClosedResourceError
  1480. elif self._transport.is_closing():
  1481. raise BrokenResourceError
  1482. else:
  1483. self._transport.sendto(*item)
  1484. class ConnectedUDPSocket(abc.ConnectedUDPSocket):
  1485. def __init__(
  1486. self, transport: asyncio.DatagramTransport, protocol: DatagramProtocol
  1487. ):
  1488. self._transport = transport
  1489. self._protocol = protocol
  1490. self._receive_guard = ResourceGuard("reading from")
  1491. self._send_guard = ResourceGuard("writing to")
  1492. self._closed = False
  1493. @property
  1494. def _raw_socket(self) -> socket.socket:
  1495. return self._transport.get_extra_info("socket")
  1496. async def aclose(self) -> None:
  1497. self._closed = True
  1498. if not self._transport.is_closing():
  1499. self._transport.close()
  1500. await self._protocol.closed_event.wait()
  1501. async def receive(self) -> bytes:
  1502. with self._receive_guard:
  1503. await AsyncIOBackend.checkpoint()
  1504. # If the buffer is empty, ask for more data
  1505. if not self._protocol.read_queue and not self._transport.is_closing():
  1506. self._protocol.read_event.clear()
  1507. await self._protocol.read_event.wait()
  1508. try:
  1509. packet = self._protocol.read_queue.popleft()
  1510. except IndexError:
  1511. if self._closed:
  1512. raise ClosedResourceError from None
  1513. else:
  1514. raise BrokenResourceError from None
  1515. return packet[0]
  1516. async def send(self, item: bytes) -> None:
  1517. with self._send_guard:
  1518. await AsyncIOBackend.checkpoint()
  1519. await self._protocol.write_event.wait()
  1520. if self._closed:
  1521. raise ClosedResourceError
  1522. elif self._transport.is_closing():
  1523. raise BrokenResourceError
  1524. else:
  1525. self._transport.sendto(item)
  1526. class UNIXDatagramSocket(_RawSocketMixin, abc.UNIXDatagramSocket):
  1527. async def receive(self) -> UNIXDatagramPacketType:
  1528. loop = get_running_loop()
  1529. await AsyncIOBackend.checkpoint()
  1530. with self._receive_guard:
  1531. while True:
  1532. try:
  1533. data = self._raw_socket.recvfrom(65536)
  1534. except BlockingIOError:
  1535. await self._wait_until_readable(loop)
  1536. except OSError as exc:
  1537. if self._closing:
  1538. raise ClosedResourceError from None
  1539. else:
  1540. raise BrokenResourceError from exc
  1541. else:
  1542. return data
  1543. async def send(self, item: UNIXDatagramPacketType) -> None:
  1544. loop = get_running_loop()
  1545. await AsyncIOBackend.checkpoint()
  1546. with self._send_guard:
  1547. while True:
  1548. try:
  1549. self._raw_socket.sendto(*item)
  1550. except BlockingIOError:
  1551. await self._wait_until_writable(loop)
  1552. except OSError as exc:
  1553. if self._closing:
  1554. raise ClosedResourceError from None
  1555. else:
  1556. raise BrokenResourceError from exc
  1557. else:
  1558. return
  1559. class ConnectedUNIXDatagramSocket(_RawSocketMixin, abc.ConnectedUNIXDatagramSocket):
  1560. async def receive(self) -> bytes:
  1561. loop = get_running_loop()
  1562. await AsyncIOBackend.checkpoint()
  1563. with self._receive_guard:
  1564. while True:
  1565. try:
  1566. data = self._raw_socket.recv(65536)
  1567. except BlockingIOError:
  1568. await self._wait_until_readable(loop)
  1569. except OSError as exc:
  1570. if self._closing:
  1571. raise ClosedResourceError from None
  1572. else:
  1573. raise BrokenResourceError from exc
  1574. else:
  1575. return data
  1576. async def send(self, item: bytes) -> None:
  1577. loop = get_running_loop()
  1578. await AsyncIOBackend.checkpoint()
  1579. with self._send_guard:
  1580. while True:
  1581. try:
  1582. self._raw_socket.send(item)
  1583. except BlockingIOError:
  1584. await self._wait_until_writable(loop)
  1585. except OSError as exc:
  1586. if self._closing:
  1587. raise ClosedResourceError from None
  1588. else:
  1589. raise BrokenResourceError from exc
  1590. else:
  1591. return
  1592. _read_events: RunVar[dict[int, asyncio.Future[bool]]] = RunVar("read_events")
  1593. _write_events: RunVar[dict[int, asyncio.Future[bool]]] = RunVar("write_events")
  1594. #
  1595. # Synchronization
  1596. #
  1597. class Event(BaseEvent):
  1598. __slots__ = ("_event",)
  1599. def __new__(cls) -> Self:
  1600. return object.__new__(cls)
  1601. def __init__(self) -> None:
  1602. self._event = asyncio.Event()
  1603. def set(self) -> None:
  1604. self._event.set()
  1605. def is_set(self) -> bool:
  1606. return self._event.is_set()
  1607. async def wait(self) -> None:
  1608. if self.is_set():
  1609. await AsyncIOBackend.checkpoint()
  1610. else:
  1611. await self._event.wait()
  1612. def statistics(self) -> EventStatistics:
  1613. return EventStatistics(len(self._event._waiters))
  1614. class Lock(BaseLock):
  1615. __slots__ = "_fast_acquire", "_owner_task", "_waiters"
  1616. def __new__(cls, *, fast_acquire: bool = False) -> Self:
  1617. return object.__new__(cls)
  1618. def __init__(self, *, fast_acquire: bool = False) -> None:
  1619. self._fast_acquire = fast_acquire
  1620. self._owner_task: asyncio.Task | None = None
  1621. self._waiters: deque[tuple[asyncio.Task, asyncio.Future]] = deque()
  1622. async def acquire(self) -> None:
  1623. task = cast(asyncio.Task, current_task())
  1624. if self._owner_task is None and not self._waiters:
  1625. await AsyncIOBackend.checkpoint_if_cancelled()
  1626. self._owner_task = task
  1627. # Unless on the "fast path", yield control of the event loop so that other
  1628. # tasks can run too
  1629. if not self._fast_acquire:
  1630. try:
  1631. await AsyncIOBackend.cancel_shielded_checkpoint()
  1632. except CancelledError:
  1633. self.release()
  1634. raise
  1635. return
  1636. if self._owner_task == task:
  1637. raise RuntimeError("Attempted to acquire an already held Lock")
  1638. fut: asyncio.Future[None] = asyncio.Future()
  1639. item = task, fut
  1640. self._waiters.append(item)
  1641. try:
  1642. await fut
  1643. except CancelledError:
  1644. if fut.cancelled():
  1645. try:
  1646. self._waiters.remove(item)
  1647. except ValueError:
  1648. pass
  1649. else:
  1650. self.release()
  1651. raise
  1652. def acquire_nowait(self) -> None:
  1653. task = cast(asyncio.Task, current_task())
  1654. if self._owner_task is None and not self._waiters:
  1655. self._owner_task = task
  1656. return
  1657. if self._owner_task is task:
  1658. raise RuntimeError("Attempted to acquire an already held Lock")
  1659. raise WouldBlock
  1660. def locked(self) -> bool:
  1661. return self._owner_task is not None
  1662. def release(self) -> None:
  1663. if self._owner_task != current_task():
  1664. raise RuntimeError("The current task is not holding this lock")
  1665. # A cancelled waiter that already received ownership removes itself from
  1666. # _waiters before calling release(); any cancelled waiter still queued here
  1667. # was cancelled before being woken, so drop it.
  1668. while self._waiters:
  1669. task, fut = self._waiters.popleft()
  1670. if fut.cancelled():
  1671. continue
  1672. self._owner_task = task
  1673. fut.set_result(None)
  1674. return
  1675. self._owner_task = None
  1676. def statistics(self) -> LockStatistics:
  1677. task_info = AsyncIOTaskInfo(self._owner_task) if self._owner_task else None
  1678. return LockStatistics(self.locked(), task_info, len(self._waiters))
  1679. class Semaphore(BaseSemaphore):
  1680. __slots__ = "_fast_acquire", "_max_value", "_value", "_waiters"
  1681. def __new__(
  1682. cls,
  1683. initial_value: int,
  1684. *,
  1685. max_value: int | None = None,
  1686. fast_acquire: bool = False,
  1687. ) -> Self:
  1688. return object.__new__(cls)
  1689. def __init__(
  1690. self,
  1691. initial_value: int,
  1692. *,
  1693. max_value: int | None = None,
  1694. fast_acquire: bool = False,
  1695. ):
  1696. super().__init__(initial_value, max_value=max_value)
  1697. self._value = initial_value
  1698. self._max_value = max_value
  1699. self._fast_acquire = fast_acquire
  1700. self._waiters: deque[asyncio.Future[None]] = deque()
  1701. async def acquire(self) -> None:
  1702. if self._value > 0 and not self._waiters:
  1703. await AsyncIOBackend.checkpoint_if_cancelled()
  1704. self._value -= 1
  1705. # Unless on the "fast path", yield control of the event loop so that other
  1706. # tasks can run too
  1707. if not self._fast_acquire:
  1708. try:
  1709. await AsyncIOBackend.cancel_shielded_checkpoint()
  1710. except CancelledError:
  1711. self.release()
  1712. raise
  1713. return
  1714. fut: asyncio.Future[None] = asyncio.Future()
  1715. self._waiters.append(fut)
  1716. try:
  1717. await fut
  1718. except CancelledError:
  1719. if fut.cancelled():
  1720. try:
  1721. self._waiters.remove(fut)
  1722. except ValueError:
  1723. pass
  1724. else:
  1725. self.release()
  1726. raise
  1727. def acquire_nowait(self) -> None:
  1728. if self._value == 0:
  1729. raise WouldBlock
  1730. self._value -= 1
  1731. def release(self) -> None:
  1732. if self._max_value is not None and self._value == self._max_value:
  1733. raise ValueError("semaphore released too many times")
  1734. while self._waiters:
  1735. fut = self._waiters.popleft()
  1736. if fut.cancelled():
  1737. continue
  1738. fut.set_result(None)
  1739. return
  1740. self._value += 1
  1741. @property
  1742. def value(self) -> int:
  1743. return self._value
  1744. @property
  1745. def max_value(self) -> int | None:
  1746. return self._max_value
  1747. def statistics(self) -> SemaphoreStatistics:
  1748. return SemaphoreStatistics(len(self._waiters))
  1749. class CapacityLimiter(BaseCapacityLimiter):
  1750. __slots__ = "_borrowers", "_total_tokens", "_wait_queue"
  1751. def __new__(cls, total_tokens: float) -> Self:
  1752. return object.__new__(cls)
  1753. def __init__(self, total_tokens: float):
  1754. self._total_tokens: float = 0
  1755. self._borrowers: set[Any] = set()
  1756. self._wait_queue: OrderedDict[Any, asyncio.Event] = OrderedDict()
  1757. self.total_tokens = total_tokens
  1758. async def __aenter__(self) -> None:
  1759. await self.acquire()
  1760. async def __aexit__(
  1761. self,
  1762. exc_type: type[BaseException] | None,
  1763. exc_val: BaseException | None,
  1764. exc_tb: TracebackType | None,
  1765. ) -> None:
  1766. self.release()
  1767. @property
  1768. def total_tokens(self) -> float:
  1769. return self._total_tokens
  1770. @total_tokens.setter
  1771. def total_tokens(self, value: float) -> None:
  1772. if not isinstance(value, int) and not math.isinf(value):
  1773. raise TypeError("total_tokens must be an int or math.inf")
  1774. if value < 0:
  1775. raise ValueError("total_tokens must be >= 0")
  1776. self._total_tokens = value
  1777. # Notify waiting tasks that they have acquired the limiter while
  1778. # there is spare capacity.
  1779. while self._wait_queue and len(self._borrowers) < self._total_tokens:
  1780. borrower, event = self._wait_queue.popitem(last=False)
  1781. self._borrowers.add(borrower)
  1782. event.set()
  1783. @property
  1784. def borrowed_tokens(self) -> int:
  1785. return len(self._borrowers)
  1786. @property
  1787. def available_tokens(self) -> float:
  1788. return self._total_tokens - len(self._borrowers)
  1789. def _notify_next_waiter(self) -> None:
  1790. """Hand a free token to the next task in line, if any."""
  1791. if self._wait_queue and len(self._borrowers) < self._total_tokens:
  1792. borrower, event = self._wait_queue.popitem(last=False)
  1793. self._borrowers.add(borrower)
  1794. event.set()
  1795. def acquire_nowait(self) -> None:
  1796. self.acquire_on_behalf_of_nowait(current_task())
  1797. def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
  1798. if borrower in self._borrowers:
  1799. raise RuntimeError(
  1800. "this borrower is already holding one of this CapacityLimiter's tokens"
  1801. )
  1802. if self._wait_queue or len(self._borrowers) >= self._total_tokens:
  1803. raise WouldBlock
  1804. self._borrowers.add(borrower)
  1805. async def acquire(self) -> None:
  1806. return await self.acquire_on_behalf_of(current_task())
  1807. async def acquire_on_behalf_of(self, borrower: object) -> None:
  1808. await AsyncIOBackend.checkpoint_if_cancelled()
  1809. try:
  1810. self.acquire_on_behalf_of_nowait(borrower)
  1811. except WouldBlock:
  1812. event = asyncio.Event()
  1813. self._wait_queue[borrower] = event
  1814. try:
  1815. await event.wait()
  1816. except BaseException:
  1817. self._wait_queue.pop(borrower, None)
  1818. if event.is_set():
  1819. self._borrowers.discard(borrower)
  1820. self._notify_next_waiter()
  1821. raise
  1822. else:
  1823. try:
  1824. await AsyncIOBackend.cancel_shielded_checkpoint()
  1825. except BaseException:
  1826. self.release()
  1827. raise
  1828. def release(self) -> None:
  1829. self.release_on_behalf_of(current_task())
  1830. def release_on_behalf_of(self, borrower: object) -> None:
  1831. try:
  1832. self._borrowers.remove(borrower)
  1833. except KeyError:
  1834. raise RuntimeError(
  1835. "this borrower isn't holding any of this CapacityLimiter's tokens"
  1836. ) from None
  1837. self._notify_next_waiter()
  1838. def statistics(self) -> CapacityLimiterStatistics:
  1839. return CapacityLimiterStatistics(
  1840. self.borrowed_tokens,
  1841. self.total_tokens,
  1842. tuple(self._borrowers),
  1843. len(self._wait_queue),
  1844. )
  1845. _default_thread_limiter: RunVar[CapacityLimiter] = RunVar("_default_thread_limiter")
  1846. #
  1847. # Operating system signals
  1848. #
  1849. class _SignalReceiver:
  1850. def __init__(self, signals: tuple[Signals, ...]):
  1851. self._signals = signals
  1852. self._loop = get_running_loop()
  1853. self._signal_queue: deque[Signals] = deque()
  1854. self._future: asyncio.Future = asyncio.Future()
  1855. self._handled_signals: set[Signals] = set()
  1856. def _deliver(self, signum: Signals) -> None:
  1857. self._signal_queue.append(signum)
  1858. if not self._future.done():
  1859. self._future.set_result(None)
  1860. def __enter__(self) -> Self:
  1861. for sig in set(self._signals):
  1862. self._loop.add_signal_handler(sig, self._deliver, sig)
  1863. self._handled_signals.add(sig)
  1864. return self
  1865. def __exit__(
  1866. self,
  1867. exc_type: type[BaseException] | None,
  1868. exc_val: BaseException | None,
  1869. exc_tb: TracebackType | None,
  1870. ) -> None:
  1871. for sig in self._handled_signals:
  1872. self._loop.remove_signal_handler(sig)
  1873. def __aiter__(self) -> _SignalReceiver:
  1874. return self
  1875. async def __anext__(self) -> Signals:
  1876. await AsyncIOBackend.checkpoint()
  1877. if not self._signal_queue:
  1878. self._future = asyncio.Future()
  1879. await self._future
  1880. return self._signal_queue.popleft()
  1881. #
  1882. # Testing and debugging
  1883. #
  1884. class AsyncIOTaskInfo(TaskInfo):
  1885. def __init__(self, task: asyncio.Task):
  1886. task_state = _task_states.get(task)
  1887. if task_state is None:
  1888. parent_id = None
  1889. else:
  1890. parent_id = task_state.parent_id
  1891. coro = task.get_coro()
  1892. assert coro is not None, "created TaskInfo from a completed Task"
  1893. super().__init__(id(task), parent_id, task.get_name(), coro)
  1894. self._task = weakref.ref(task)
  1895. def has_pending_cancellation(self) -> bool:
  1896. if not (task := self._task()):
  1897. # If the task isn't around anymore, it won't have a pending cancellation
  1898. return False
  1899. if task._must_cancel or ( # type: ignore[attr-defined]
  1900. isinstance(task._fut_waiter, asyncio.Future) # type: ignore[attr-defined]
  1901. and task._fut_waiter.cancelled() # type: ignore[attr-defined]
  1902. ):
  1903. return True
  1904. if task_state := _task_states.get(task):
  1905. if cancel_scope := task_state.cancel_scope:
  1906. return cancel_scope._effectively_cancelled
  1907. return False
  1908. class TestRunner(abc.TestRunner):
  1909. _send_stream: MemoryObjectSendStream[tuple[Awaitable[Any], asyncio.Future[Any]]]
  1910. def __init__(
  1911. self,
  1912. *,
  1913. debug: bool | None = None,
  1914. use_uvloop: bool = False,
  1915. loop_factory: Callable[[], AbstractEventLoop] | None = None,
  1916. ) -> None:
  1917. if use_uvloop and loop_factory is None:
  1918. if sys.platform != "win32":
  1919. import uvloop
  1920. loop_factory = uvloop.new_event_loop
  1921. else:
  1922. import winloop
  1923. loop_factory = winloop.new_event_loop
  1924. self._runner = Runner(debug=debug, loop_factory=loop_factory)
  1925. self._exceptions: list[BaseException] = []
  1926. self._runner_task: asyncio.Task | None = None
  1927. def __enter__(self) -> Self:
  1928. self._runner.__enter__()
  1929. self.get_loop().set_exception_handler(self._exception_handler)
  1930. return self
  1931. def __exit__(
  1932. self,
  1933. exc_type: type[BaseException] | None,
  1934. exc_val: BaseException | None,
  1935. exc_tb: TracebackType | None,
  1936. ) -> None:
  1937. self._runner.__exit__(exc_type, exc_val, exc_tb)
  1938. def get_loop(self) -> AbstractEventLoop:
  1939. return self._runner.get_loop()
  1940. def is_running(self) -> bool:
  1941. try:
  1942. asyncio.get_running_loop()
  1943. return True
  1944. except RuntimeError:
  1945. return False
  1946. def _exception_handler(
  1947. self, loop: asyncio.AbstractEventLoop, context: dict[str, Any]
  1948. ) -> None:
  1949. if isinstance(context.get("exception"), Exception):
  1950. self._exceptions.append(context["exception"])
  1951. else:
  1952. loop.default_exception_handler(context)
  1953. def _raise_async_exceptions(self) -> None:
  1954. # Re-raise any exceptions raised in asynchronous callbacks
  1955. if self._exceptions:
  1956. exceptions, self._exceptions = self._exceptions, []
  1957. if len(exceptions) == 1:
  1958. raise exceptions[0]
  1959. elif exceptions:
  1960. raise BaseExceptionGroup(
  1961. "Multiple exceptions occurred in asynchronous callbacks", exceptions
  1962. )
  1963. async def _run_tests_and_fixtures(
  1964. self,
  1965. receive_stream: MemoryObjectReceiveStream[
  1966. tuple[Awaitable[T_Retval], asyncio.Future[T_Retval]]
  1967. ],
  1968. ) -> None:
  1969. from _pytest.outcomes import OutcomeException
  1970. with receive_stream, self._send_stream:
  1971. async for coro, future in receive_stream:
  1972. try:
  1973. retval = await coro
  1974. except CancelledError as exc:
  1975. if not future.cancelled():
  1976. future.cancel(*exc.args)
  1977. raise
  1978. except BaseException as exc:
  1979. if not future.cancelled():
  1980. future.set_exception(exc)
  1981. if not isinstance(exc, (Exception, OutcomeException)):
  1982. raise
  1983. else:
  1984. if not future.cancelled():
  1985. future.set_result(retval)
  1986. async def _call_in_runner_task(
  1987. self,
  1988. func: Callable[P, Awaitable[T_Retval]],
  1989. /,
  1990. *args: P.args,
  1991. **kwargs: P.kwargs,
  1992. ) -> T_Retval:
  1993. if not self._runner_task:
  1994. self._send_stream, receive_stream = create_memory_object_stream[
  1995. tuple[Awaitable[Any], asyncio.Future]
  1996. ](1)
  1997. self._runner_task = self.get_loop().create_task(
  1998. self._run_tests_and_fixtures(receive_stream)
  1999. )
  2000. coro = func(*args, **kwargs)
  2001. future: asyncio.Future[T_Retval] = self.get_loop().create_future()
  2002. self._send_stream.send_nowait((coro, future))
  2003. return await future
  2004. def run_asyncgen_fixture(
  2005. self,
  2006. fixture_func: Callable[..., AsyncGenerator[T_Retval, Any]],
  2007. kwargs: dict[str, Any],
  2008. ) -> Iterable[T_Retval]:
  2009. asyncgen = fixture_func(**kwargs)
  2010. fixturevalue: T_Retval = self.get_loop().run_until_complete(
  2011. self._call_in_runner_task(asyncgen.asend, None)
  2012. )
  2013. self._raise_async_exceptions()
  2014. yield fixturevalue
  2015. try:
  2016. self.get_loop().run_until_complete(
  2017. self._call_in_runner_task(asyncgen.asend, None)
  2018. )
  2019. except StopAsyncIteration:
  2020. self._raise_async_exceptions()
  2021. else:
  2022. self.get_loop().run_until_complete(asyncgen.aclose())
  2023. raise RuntimeError("Async generator fixture did not stop")
  2024. def run_fixture(
  2025. self,
  2026. fixture_func: Callable[..., Coroutine[Any, Any, T_Retval]],
  2027. kwargs: dict[str, Any],
  2028. ) -> T_Retval:
  2029. retval = self.get_loop().run_until_complete(
  2030. self._call_in_runner_task(fixture_func, **kwargs)
  2031. )
  2032. self._raise_async_exceptions()
  2033. return retval
  2034. def run_test(
  2035. self, test_func: Callable[..., Coroutine[Any, Any, Any]], kwargs: dict[str, Any]
  2036. ) -> None:
  2037. from _pytest.outcomes import OutcomeException
  2038. try:
  2039. self.get_loop().run_until_complete(
  2040. self._call_in_runner_task(test_func, **kwargs)
  2041. )
  2042. except Exception as exc:
  2043. self._exceptions.append(exc)
  2044. except OutcomeException:
  2045. raise
  2046. except BaseException:
  2047. # A BaseException (e.g. KeyboardInterrupt, SystemExit) interrupted the event loop before
  2048. # the test completed. Cancel _runner_task so it does not resume when the event
  2049. # loop is re-entered during async generator fixture teardown.
  2050. if self._runner_task is not None and not self._runner_task.done():
  2051. self._runner_task.cancel()
  2052. self._send_stream.close()
  2053. try:
  2054. self.get_loop().run_until_complete(self._runner_task)
  2055. except CancelledError:
  2056. pass
  2057. finally:
  2058. self._runner_task = None
  2059. raise
  2060. self._raise_async_exceptions()
  2061. class _ProcessStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
  2062. """
  2063. A subprocess protocol that allows us to be notified of ``process_exited``
  2064. asyncio's own ``Process.wait()`` only resolves once every pipe transport has
  2065. disconnected so to get same semantics as on trio and uvloop we need this.
  2066. """
  2067. def __init__(self) -> None:
  2068. # Match the standard factory for asyncio.create_process
  2069. super().__init__(limit=2**16, loop=asyncio.get_running_loop())
  2070. self.exited = asyncio.Event()
  2071. def process_exited(self) -> None:
  2072. super().process_exited()
  2073. self.exited.set()
  2074. class AsyncIOBackend(AsyncBackend):
  2075. @classmethod
  2076. def run(
  2077. cls,
  2078. func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval]],
  2079. args: tuple[Unpack[PosArgsT]],
  2080. kwargs: dict[str, Any],
  2081. options: dict[str, Any],
  2082. ) -> T_Retval:
  2083. @wraps(func)
  2084. async def wrapper() -> T_Retval:
  2085. task = cast(asyncio.Task, current_task())
  2086. task.set_name(get_callable_name(func))
  2087. _task_states[task] = TaskState(None, None)
  2088. try:
  2089. return await func(*args)
  2090. finally:
  2091. del _task_states[task]
  2092. debug = options.get("debug", None)
  2093. loop_factory = options.get("loop_factory", None)
  2094. if loop_factory is None and options.get("use_uvloop", False):
  2095. if sys.platform != "win32":
  2096. import uvloop
  2097. loop_factory = uvloop.new_event_loop
  2098. else:
  2099. import winloop
  2100. loop_factory = winloop.new_event_loop
  2101. with Runner(debug=debug, loop_factory=loop_factory) as runner:
  2102. return runner.run(wrapper())
  2103. @classmethod
  2104. def current_token(cls) -> object:
  2105. return get_running_loop()
  2106. @classmethod
  2107. def current_time(cls) -> float:
  2108. return get_running_loop().time()
  2109. @classmethod
  2110. def cancelled_exception_class(cls) -> type[BaseException]:
  2111. return CancelledError
  2112. @classmethod
  2113. async def checkpoint(cls) -> None:
  2114. await sleep(0)
  2115. @classmethod
  2116. async def checkpoint_if_cancelled(cls) -> None:
  2117. task = current_task()
  2118. if task is None:
  2119. return
  2120. try:
  2121. cancel_scope = _task_states[task].cancel_scope
  2122. except KeyError:
  2123. return
  2124. while cancel_scope:
  2125. if cancel_scope.cancel_called:
  2126. await sleep(0)
  2127. elif cancel_scope.shield:
  2128. break
  2129. else:
  2130. cancel_scope = cancel_scope._parent_scope
  2131. @classmethod
  2132. async def cancel_shielded_checkpoint(cls) -> None:
  2133. with CancelScope(shield=True):
  2134. await sleep(0)
  2135. @classmethod
  2136. async def sleep(cls, delay: float) -> None:
  2137. await sleep(delay)
  2138. @classmethod
  2139. def create_cancel_scope(
  2140. cls, *, deadline: float = math.inf, shield: bool = False
  2141. ) -> CancelScope:
  2142. return CancelScope(deadline=deadline, shield=shield)
  2143. @classmethod
  2144. def current_effective_deadline(cls) -> float:
  2145. if (task := current_task()) is None:
  2146. return math.inf
  2147. try:
  2148. cancel_scope = _task_states[task].cancel_scope
  2149. except KeyError:
  2150. return math.inf
  2151. deadline = math.inf
  2152. while cancel_scope:
  2153. deadline = min(deadline, cancel_scope.deadline)
  2154. if cancel_scope._cancel_called:
  2155. deadline = -math.inf
  2156. break
  2157. elif cancel_scope.shield:
  2158. break
  2159. else:
  2160. cancel_scope = cancel_scope._parent_scope
  2161. return deadline
  2162. @classmethod
  2163. def create_task_group(cls) -> abc.TaskGroup:
  2164. return TaskGroup()
  2165. @classmethod
  2166. def create_event(cls) -> BaseEvent:
  2167. return Event()
  2168. @classmethod
  2169. def create_lock(cls, *, fast_acquire: bool) -> BaseLock:
  2170. return Lock(fast_acquire=fast_acquire)
  2171. @classmethod
  2172. def create_semaphore(
  2173. cls,
  2174. initial_value: int,
  2175. *,
  2176. max_value: int | None = None,
  2177. fast_acquire: bool = False,
  2178. ) -> BaseSemaphore:
  2179. return Semaphore(initial_value, max_value=max_value, fast_acquire=fast_acquire)
  2180. @classmethod
  2181. def create_capacity_limiter(cls, total_tokens: float) -> BaseCapacityLimiter:
  2182. return CapacityLimiter(total_tokens)
  2183. @classmethod
  2184. async def run_sync_in_worker_thread( # type: ignore[return]
  2185. cls,
  2186. func: Callable[[Unpack[PosArgsT]], T_Retval],
  2187. args: tuple[Unpack[PosArgsT]],
  2188. abandon_on_cancel: bool = False,
  2189. limiter: BaseCapacityLimiter | None = None,
  2190. ) -> T_Retval:
  2191. await cls.checkpoint()
  2192. # If this is the first run in this event loop thread, set up the necessary
  2193. # variables
  2194. try:
  2195. idle_workers = _threadpool_idle_workers.get()
  2196. workers = _threadpool_workers.get()
  2197. except LookupError:
  2198. idle_workers = deque()
  2199. workers = set()
  2200. _threadpool_idle_workers.set(idle_workers)
  2201. _threadpool_workers.set(workers)
  2202. async with limiter or cls.current_default_thread_limiter():
  2203. with CancelScope(shield=not abandon_on_cancel) as scope:
  2204. future = asyncio.Future[T_Retval]()
  2205. root_task = find_root_task()
  2206. if not idle_workers:
  2207. worker = WorkerThread(root_task, workers, idle_workers)
  2208. worker.start()
  2209. workers.add(worker)
  2210. root_task.add_done_callback(worker.stop, context=Context())
  2211. else:
  2212. worker = idle_workers.pop()
  2213. # Prune any other workers that have been idle for MAX_IDLE_TIME
  2214. # seconds or longer
  2215. now = cls.current_time()
  2216. while idle_workers:
  2217. if (
  2218. now - idle_workers[0].idle_since
  2219. < WorkerThread.MAX_IDLE_TIME
  2220. ):
  2221. break
  2222. expired_worker = idle_workers.popleft()
  2223. expired_worker.root_task.remove_done_callback(
  2224. expired_worker.stop
  2225. )
  2226. expired_worker.stop()
  2227. context = copy_context()
  2228. context.run(set_current_async_library, None)
  2229. if abandon_on_cancel or scope._parent_scope is None:
  2230. worker_scope = scope
  2231. else:
  2232. worker_scope = scope._parent_scope
  2233. worker.queue.put_nowait((context, func, args, future, worker_scope))
  2234. return await future
  2235. @classmethod
  2236. def check_cancelled(cls) -> None:
  2237. scope: CancelScope | None = threadlocals.current_cancel_scope
  2238. while scope is not None:
  2239. if scope.cancel_called:
  2240. raise CancelledError(f"Cancelled via cancel scope {id(scope):x}")
  2241. if scope.shield:
  2242. return
  2243. scope = scope._parent_scope
  2244. @classmethod
  2245. def run_async_from_thread(
  2246. cls,
  2247. func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
  2248. args: tuple[Unpack[PosArgsT]],
  2249. token: object,
  2250. ) -> T_co:
  2251. async def task_wrapper() -> T_co:
  2252. __tracebackhide__ = True
  2253. if scope is not None:
  2254. task = cast(asyncio.Task, current_task())
  2255. _task_states[task] = TaskState(None, scope)
  2256. scope._tasks.add(task)
  2257. try:
  2258. return await func(*args)
  2259. except CancelledError as exc:
  2260. raise concurrent.futures.CancelledError(str(exc)) from None
  2261. finally:
  2262. if scope is not None:
  2263. scope._tasks.discard(task)
  2264. loop = cast(
  2265. "AbstractEventLoop", token or threadlocals.current_token.native_token
  2266. )
  2267. if loop.is_closed():
  2268. raise RunFinishedError
  2269. context = copy_context()
  2270. context.run(set_current_async_library, "asyncio")
  2271. scope = getattr(threadlocals, "current_cancel_scope", None)
  2272. f: concurrent.futures.Future[T_co] = context.run(
  2273. asyncio.run_coroutine_threadsafe, task_wrapper(), loop=loop
  2274. )
  2275. return f.result()
  2276. @classmethod
  2277. def run_sync_from_thread(
  2278. cls,
  2279. func: Callable[[Unpack[PosArgsT]], T_Retval],
  2280. args: tuple[Unpack[PosArgsT]],
  2281. token: object,
  2282. ) -> T_Retval:
  2283. @wraps(func)
  2284. def wrapper() -> None:
  2285. try:
  2286. set_current_async_library("asyncio")
  2287. f.set_result(func(*args))
  2288. except BaseException as exc:
  2289. f.set_exception(exc)
  2290. if not isinstance(exc, Exception):
  2291. raise
  2292. loop = cast(
  2293. "AbstractEventLoop", token or threadlocals.current_token.native_token
  2294. )
  2295. if loop.is_closed():
  2296. raise RunFinishedError
  2297. f: concurrent.futures.Future[T_Retval] = Future()
  2298. loop.call_soon_threadsafe(wrapper)
  2299. return f.result()
  2300. @classmethod
  2301. async def open_process(
  2302. cls,
  2303. command: StrOrBytesPath | Sequence[StrOrBytesPath],
  2304. *,
  2305. stdin: int | IO[Any] | None,
  2306. stdout: int | IO[Any] | None,
  2307. stderr: int | IO[Any] | None,
  2308. **kwargs: Any,
  2309. ) -> Process:
  2310. await cls.checkpoint()
  2311. if isinstance(command, PathLike):
  2312. command = os.fspath(command)
  2313. # Use loop.subprocess_shell()/subprocess_exec() rather than their
  2314. # asyncio.create_subprocess_*() counterparts to get access to
  2315. # transport/protocol.
  2316. loop = asyncio.get_running_loop()
  2317. if isinstance(command, (str, bytes)):
  2318. transport, protocol = await loop.subprocess_shell(
  2319. _ProcessStreamProtocol,
  2320. command,
  2321. stdin=stdin,
  2322. stdout=stdout,
  2323. stderr=stderr,
  2324. **kwargs,
  2325. )
  2326. else:
  2327. transport, protocol = await loop.subprocess_exec(
  2328. _ProcessStreamProtocol,
  2329. *command,
  2330. stdin=stdin,
  2331. stdout=stdout,
  2332. stderr=stderr,
  2333. **kwargs,
  2334. )
  2335. process = asyncio.subprocess.Process(transport, protocol, loop)
  2336. stdin_stream = StreamWriterWrapper(process.stdin) if process.stdin else None
  2337. stdout_stream = StreamReaderWrapper(process.stdout) if process.stdout else None
  2338. stderr_stream = StreamReaderWrapper(process.stderr) if process.stderr else None
  2339. return Process(
  2340. process,
  2341. stdin_stream,
  2342. stdout_stream,
  2343. stderr_stream,
  2344. protocol.exited,
  2345. transport,
  2346. )
  2347. @classmethod
  2348. def setup_process_pool_exit_at_shutdown(cls, workers: set[abc.Process]) -> None:
  2349. create_task(
  2350. _shutdown_process_pool_on_exit(workers),
  2351. name="AnyIO process pool shutdown task",
  2352. )
  2353. find_root_task().add_done_callback(
  2354. partial(_forcibly_shutdown_process_pool_on_exit, workers) # type:ignore[arg-type]
  2355. )
  2356. @classmethod
  2357. async def connect_tcp(
  2358. cls, host: str, port: int, local_address: IPSockAddrType | None = None
  2359. ) -> abc.SocketStream:
  2360. transport, protocol = cast(
  2361. tuple[asyncio.Transport, StreamProtocol],
  2362. await get_running_loop().create_connection(
  2363. StreamProtocol, host, port, local_addr=local_address
  2364. ),
  2365. )
  2366. transport.pause_reading()
  2367. return SocketStream(transport, protocol)
  2368. @classmethod
  2369. async def connect_unix(cls, path: str | bytes) -> abc.UNIXSocketStream:
  2370. await cls.checkpoint()
  2371. loop = get_running_loop()
  2372. raw_socket = socket.socket(socket.AF_UNIX)
  2373. raw_socket.setblocking(False)
  2374. while True:
  2375. try:
  2376. raw_socket.connect(path)
  2377. except BlockingIOError:
  2378. f: asyncio.Future = asyncio.Future()
  2379. loop.add_writer(raw_socket, f.set_result, None)
  2380. f.add_done_callback(lambda _: loop.remove_writer(raw_socket))
  2381. await f
  2382. except BaseException:
  2383. raw_socket.close()
  2384. raise
  2385. else:
  2386. return UNIXSocketStream(raw_socket)
  2387. @classmethod
  2388. def create_tcp_listener(cls, sock: socket.socket) -> SocketListener:
  2389. return TCPSocketListener(sock)
  2390. @classmethod
  2391. def create_unix_listener(cls, sock: socket.socket) -> SocketListener:
  2392. return UNIXSocketListener(sock)
  2393. @classmethod
  2394. async def create_udp_socket(
  2395. cls,
  2396. family: AddressFamily,
  2397. local_address: IPSockAddrType | None,
  2398. remote_address: IPSockAddrType | None,
  2399. reuse_port: bool,
  2400. ) -> UDPSocket | ConnectedUDPSocket:
  2401. transport, protocol = await get_running_loop().create_datagram_endpoint(
  2402. DatagramProtocol,
  2403. local_addr=local_address,
  2404. remote_addr=remote_address,
  2405. family=family,
  2406. reuse_port=reuse_port,
  2407. )
  2408. if protocol.exception:
  2409. transport.close()
  2410. raise protocol.exception
  2411. if not remote_address:
  2412. return UDPSocket(transport, protocol)
  2413. else:
  2414. return ConnectedUDPSocket(transport, protocol)
  2415. @classmethod
  2416. async def create_unix_datagram_socket( # type: ignore[override]
  2417. cls, raw_socket: socket.socket, remote_path: str | bytes | None
  2418. ) -> abc.UNIXDatagramSocket | abc.ConnectedUNIXDatagramSocket:
  2419. await cls.checkpoint()
  2420. loop = get_running_loop()
  2421. if remote_path:
  2422. while True:
  2423. try:
  2424. raw_socket.connect(remote_path)
  2425. except BlockingIOError:
  2426. f: asyncio.Future = asyncio.Future()
  2427. loop.add_writer(raw_socket, f.set_result, None)
  2428. f.add_done_callback(lambda _: loop.remove_writer(raw_socket))
  2429. await f
  2430. except BaseException:
  2431. raw_socket.close()
  2432. raise
  2433. else:
  2434. return ConnectedUNIXDatagramSocket(raw_socket)
  2435. else:
  2436. return UNIXDatagramSocket(raw_socket)
  2437. @classmethod
  2438. async def getaddrinfo(
  2439. cls,
  2440. host: bytes | str | None,
  2441. port: str | int | None,
  2442. *,
  2443. family: int | AddressFamily = 0,
  2444. type: int | SocketKind = 0,
  2445. proto: int = 0,
  2446. flags: int = 0,
  2447. ) -> Sequence[
  2448. tuple[
  2449. AddressFamily,
  2450. SocketKind,
  2451. int,
  2452. str,
  2453. tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes],
  2454. ]
  2455. ]:
  2456. return await get_running_loop().getaddrinfo(
  2457. host, port, family=family, type=type, proto=proto, flags=flags
  2458. )
  2459. @classmethod
  2460. async def getnameinfo(
  2461. cls, sockaddr: IPSockAddrType, flags: int = 0
  2462. ) -> tuple[str, str]:
  2463. return await get_running_loop().getnameinfo(sockaddr, flags)
  2464. @classmethod
  2465. async def wait_readable(cls, obj: FileDescriptorLike) -> None:
  2466. try:
  2467. read_events = _read_events.get()
  2468. except LookupError:
  2469. read_events = {}
  2470. _read_events.set(read_events)
  2471. fd = obj if isinstance(obj, int) else obj.fileno()
  2472. if read_events.get(fd):
  2473. raise BusyResourceError("reading from")
  2474. loop = get_running_loop()
  2475. fut: asyncio.Future[bool] = loop.create_future()
  2476. def cb() -> None:
  2477. try:
  2478. del read_events[fd]
  2479. except KeyError:
  2480. pass
  2481. else:
  2482. remove_reader(fd)
  2483. try:
  2484. fut.set_result(True)
  2485. except asyncio.InvalidStateError:
  2486. pass
  2487. try:
  2488. loop.add_reader(fd, cb)
  2489. except NotImplementedError:
  2490. from anyio._core._asyncio_selector_thread import get_selector
  2491. selector = get_selector()
  2492. selector.add_reader(fd, cb)
  2493. remove_reader = selector.remove_reader
  2494. else:
  2495. remove_reader = loop.remove_reader
  2496. read_events[fd] = fut
  2497. try:
  2498. success = await fut
  2499. finally:
  2500. try:
  2501. del read_events[fd]
  2502. except KeyError:
  2503. pass
  2504. else:
  2505. remove_reader(fd)
  2506. if not success:
  2507. raise ClosedResourceError
  2508. @classmethod
  2509. async def wait_writable(cls, obj: FileDescriptorLike) -> None:
  2510. try:
  2511. write_events = _write_events.get()
  2512. except LookupError:
  2513. write_events = {}
  2514. _write_events.set(write_events)
  2515. fd = obj if isinstance(obj, int) else obj.fileno()
  2516. if write_events.get(fd):
  2517. raise BusyResourceError("writing to")
  2518. loop = get_running_loop()
  2519. fut: asyncio.Future[bool] = loop.create_future()
  2520. def cb() -> None:
  2521. try:
  2522. del write_events[fd]
  2523. except KeyError:
  2524. pass
  2525. else:
  2526. remove_writer(fd)
  2527. try:
  2528. fut.set_result(True)
  2529. except asyncio.InvalidStateError:
  2530. pass
  2531. try:
  2532. loop.add_writer(fd, cb)
  2533. except NotImplementedError:
  2534. from anyio._core._asyncio_selector_thread import get_selector
  2535. selector = get_selector()
  2536. selector.add_writer(fd, cb)
  2537. remove_writer = selector.remove_writer
  2538. else:
  2539. remove_writer = loop.remove_writer
  2540. write_events[fd] = fut
  2541. try:
  2542. success = await fut
  2543. finally:
  2544. try:
  2545. del write_events[fd]
  2546. except KeyError:
  2547. pass
  2548. else:
  2549. remove_writer(fd)
  2550. if not success:
  2551. raise ClosedResourceError
  2552. @classmethod
  2553. def notify_closing(cls, obj: FileDescriptorLike) -> None:
  2554. fd = obj if isinstance(obj, int) else obj.fileno()
  2555. loop = get_running_loop()
  2556. try:
  2557. write_events = _write_events.get()
  2558. except LookupError:
  2559. pass
  2560. else:
  2561. try:
  2562. fut = write_events.pop(fd)
  2563. except KeyError:
  2564. pass
  2565. else:
  2566. try:
  2567. fut.set_result(False)
  2568. except asyncio.InvalidStateError:
  2569. pass
  2570. try:
  2571. loop.remove_writer(fd)
  2572. except NotImplementedError:
  2573. from anyio._core._asyncio_selector_thread import get_selector
  2574. get_selector().remove_writer(fd)
  2575. try:
  2576. read_events = _read_events.get()
  2577. except LookupError:
  2578. pass
  2579. else:
  2580. try:
  2581. fut = read_events.pop(fd)
  2582. except KeyError:
  2583. pass
  2584. else:
  2585. try:
  2586. fut.set_result(False)
  2587. except asyncio.InvalidStateError:
  2588. pass
  2589. try:
  2590. loop.remove_reader(fd)
  2591. except NotImplementedError:
  2592. from anyio._core._asyncio_selector_thread import get_selector
  2593. get_selector().remove_reader(fd)
  2594. @classmethod
  2595. async def wrap_listener_socket(cls, sock: socket.socket) -> SocketListener:
  2596. if hasattr(socket, "AF_UNIX") and sock.family == socket.AF_UNIX:
  2597. return UNIXSocketListener(sock)
  2598. return TCPSocketListener(sock)
  2599. @classmethod
  2600. async def wrap_stream_socket(cls, sock: socket.socket) -> SocketStream:
  2601. transport, protocol = await get_running_loop().create_connection(
  2602. StreamProtocol, sock=sock
  2603. )
  2604. return SocketStream(transport, protocol)
  2605. @classmethod
  2606. async def wrap_unix_stream_socket(cls, sock: socket.socket) -> UNIXSocketStream:
  2607. return UNIXSocketStream(sock)
  2608. @classmethod
  2609. async def wrap_udp_socket(cls, sock: socket.socket) -> UDPSocket:
  2610. transport, protocol = await get_running_loop().create_datagram_endpoint(
  2611. DatagramProtocol, sock=sock
  2612. )
  2613. return UDPSocket(transport, protocol)
  2614. @classmethod
  2615. async def wrap_connected_udp_socket(cls, sock: socket.socket) -> ConnectedUDPSocket:
  2616. transport, protocol = await get_running_loop().create_datagram_endpoint(
  2617. DatagramProtocol, sock=sock
  2618. )
  2619. return ConnectedUDPSocket(transport, protocol)
  2620. @classmethod
  2621. async def wrap_unix_datagram_socket(cls, sock: socket.socket) -> UNIXDatagramSocket:
  2622. return UNIXDatagramSocket(sock)
  2623. @classmethod
  2624. async def wrap_connected_unix_datagram_socket(
  2625. cls, sock: socket.socket
  2626. ) -> ConnectedUNIXDatagramSocket:
  2627. return ConnectedUNIXDatagramSocket(sock)
  2628. @classmethod
  2629. def current_default_thread_limiter(cls) -> CapacityLimiter:
  2630. try:
  2631. return _default_thread_limiter.get()
  2632. except LookupError:
  2633. limiter = CapacityLimiter(40)
  2634. _default_thread_limiter.set(limiter)
  2635. return limiter
  2636. @classmethod
  2637. def open_signal_receiver(
  2638. cls, *signals: Signals
  2639. ) -> AbstractContextManager[AsyncIterator[Signals]]:
  2640. return _SignalReceiver(signals)
  2641. @classmethod
  2642. def get_current_task(cls) -> TaskInfo:
  2643. return AsyncIOTaskInfo(current_task()) # type: ignore[arg-type]
  2644. @classmethod
  2645. def get_running_tasks(cls) -> Sequence[TaskInfo]:
  2646. return [AsyncIOTaskInfo(task) for task in all_tasks() if not task.done()]
  2647. @classmethod
  2648. async def wait_all_tasks_blocked(cls) -> None:
  2649. await cls.checkpoint()
  2650. this_task = current_task()
  2651. while True:
  2652. for task in all_tasks():
  2653. if task is this_task:
  2654. continue
  2655. waiter = task._fut_waiter # type: ignore[attr-defined]
  2656. if waiter is None or waiter.done():
  2657. await sleep(0.1)
  2658. break
  2659. else:
  2660. return
  2661. @classmethod
  2662. def create_test_runner(cls, options: dict[str, Any]) -> TestRunner:
  2663. return TestRunner(**options)
  2664. backend_class = AsyncIOBackend