_asyncio.py 97 KB

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