_trio.py 45 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473
  1. from __future__ import annotations
  2. import array
  3. import math
  4. import os
  5. import socket
  6. import sys
  7. import types
  8. import weakref
  9. from collections.abc import (
  10. AsyncGenerator,
  11. AsyncIterator,
  12. Awaitable,
  13. Callable,
  14. Collection,
  15. Coroutine,
  16. Iterable,
  17. Sequence,
  18. )
  19. from contextlib import AbstractContextManager
  20. from contextvars import Context
  21. from dataclasses import dataclass
  22. from functools import partial, wraps
  23. from io import IOBase
  24. from os import PathLike
  25. from signal import Signals
  26. from socket import AddressFamily, SocketKind
  27. from types import TracebackType
  28. from typing import (
  29. IO,
  30. TYPE_CHECKING,
  31. Any,
  32. Generic,
  33. Literal,
  34. NoReturn,
  35. ParamSpec,
  36. TypeVar,
  37. cast,
  38. overload,
  39. )
  40. import trio.from_thread
  41. import trio.lowlevel
  42. from outcome import Error, Outcome, Value
  43. from trio.lowlevel import (
  44. current_root_task,
  45. current_task,
  46. notify_closing,
  47. wait_readable,
  48. wait_writable,
  49. )
  50. from trio.socket import SocketType as TrioSocketType
  51. from trio.to_thread import run_sync
  52. from .. import (
  53. CapacityLimiterStatistics,
  54. EventStatistics,
  55. LockStatistics,
  56. RunFinishedError,
  57. TaskInfo,
  58. WouldBlock,
  59. abc,
  60. )
  61. from .._core._eventloop import claim_worker_thread
  62. from .._core._exceptions import (
  63. BrokenResourceError,
  64. BusyResourceError,
  65. ClosedResourceError,
  66. EndOfStream,
  67. )
  68. from .._core._sockets import convert_ipv6_sockaddr
  69. from .._core._streams import create_memory_object_stream
  70. from .._core._synchronization import (
  71. CapacityLimiter as BaseCapacityLimiter,
  72. )
  73. from .._core._synchronization import Event as BaseEvent
  74. from .._core._synchronization import Lock as BaseLock
  75. from .._core._synchronization import (
  76. ResourceGuard,
  77. SemaphoreStatistics,
  78. )
  79. from .._core._synchronization import Semaphore as BaseSemaphore
  80. from .._core._tasks import CancelScope as BaseCancelScope
  81. from .._core._tasks import TaskHandle
  82. from ..abc import IPSockAddrType, UDPPacketType, UNIXDatagramPacketType
  83. from ..abc._eventloop import AsyncBackend, StrOrBytesPath
  84. from ..abc._tasks import T_contra, call_for_coroutine, get_callable_name, get_coro_name
  85. if TYPE_CHECKING:
  86. from _typeshed import FileDescriptorLike
  87. from ..streams.memory import MemoryObjectSendStream
  88. if sys.version_info < (3, 15):
  89. from typing_extensions import sentinel
  90. if sys.version_info >= (3, 11):
  91. from typing import Self, TypeVarTuple, Unpack
  92. else:
  93. from exceptiongroup import BaseExceptionGroup
  94. from typing_extensions import Self, TypeVarTuple, Unpack
  95. T = TypeVar("T")
  96. T_Retval = TypeVar("T_Retval")
  97. T_co = TypeVar("T_co", covariant=True)
  98. T_SockAddr = TypeVar("T_SockAddr", str, IPSockAddrType)
  99. PosArgsT = TypeVarTuple("PosArgsT")
  100. P = ParamSpec("P")
  101. def ensure_returns_coro(
  102. func: Callable[P, Awaitable[T_Retval]],
  103. ) -> Callable[P, Coroutine[Any, Any, T_Retval]]:
  104. @wraps(func)
  105. def wrapper(*args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, T_Retval]:
  106. awaitable = func(*args, **kwargs)
  107. # Check the common case first.
  108. if isinstance(awaitable, Coroutine):
  109. return awaitable
  110. elif not isinstance(awaitable, Awaitable):
  111. # The user violated the type annotations. Still, we should pass this on to
  112. # Trio so it can raise with an appropriate message.
  113. return awaitable
  114. else:
  115. @wraps(func)
  116. async def inner_wrapper() -> T_Retval:
  117. return await awaitable
  118. return inner_wrapper()
  119. return wrapper
  120. #
  121. # Event loop
  122. #
  123. RunVar = trio.lowlevel.RunVar
  124. #
  125. # Timeouts and cancellation
  126. #
  127. class CancelScope(BaseCancelScope):
  128. __slots__ = ("__original",)
  129. def __new__(
  130. cls, original: trio.CancelScope | None = None, **kwargs: object
  131. ) -> Self:
  132. return object.__new__(cls)
  133. def __init__(self, original: trio.CancelScope | None = None, **kwargs: Any) -> None:
  134. self.__original = original or trio.CancelScope(**kwargs)
  135. def __enter__(self) -> Self:
  136. self.__original.__enter__()
  137. return self
  138. def __exit__(
  139. self,
  140. exc_type: type[BaseException] | None,
  141. exc_val: BaseException | None,
  142. exc_tb: TracebackType | None,
  143. ) -> bool:
  144. return self.__original.__exit__(exc_type, exc_val, exc_tb)
  145. def cancel(self, reason: str | None = None) -> None:
  146. self.__original.cancel(reason)
  147. @property
  148. def deadline(self) -> float:
  149. return self.__original.deadline
  150. @deadline.setter
  151. def deadline(self, value: float) -> None:
  152. self.__original.deadline = value
  153. @property
  154. def cancel_called(self) -> bool:
  155. return self.__original.cancel_called
  156. @property
  157. def cancelled_caught(self) -> bool:
  158. return self.__original.cancelled_caught
  159. @property
  160. def shield(self) -> bool:
  161. return self.__original.shield
  162. @shield.setter
  163. def shield(self, value: bool) -> None:
  164. self.__original.shield = value
  165. #
  166. # Task groups
  167. #
  168. empty_start_value = sentinel("empty_start_value")
  169. class _TrioTaskStatus(abc.TaskStatus[T_contra], Generic[T_contra]):
  170. early_start_value: T_contra | object = empty_start_value
  171. real_task_status: trio.TaskStatus[T_contra | None] | None = None
  172. def started(self, value: T_contra | None = None) -> None:
  173. if self.real_task_status is None:
  174. if self.early_start_value is not empty_start_value:
  175. raise RuntimeError("called 'started' twice on the same task status")
  176. self.early_start_value = value
  177. else:
  178. self.real_task_status.started(value)
  179. class TaskGroup(abc.TaskGroup):
  180. def __init__(self) -> None:
  181. self._entered = False
  182. self._active = False
  183. self._nursery_manager = trio.open_nursery(strict_exception_groups=True)
  184. self.cancel_scope = None # type: ignore[assignment]
  185. async def __aenter__(self) -> Self:
  186. if self._entered:
  187. raise RuntimeError("TaskGroup cannot be entered more than once")
  188. self._entered = True
  189. self._active = True
  190. self._nursery = await self._nursery_manager.__aenter__()
  191. self.cancel_scope = CancelScope(self._nursery.cancel_scope)
  192. return self
  193. async def __aexit__(
  194. self,
  195. exc_type: type[BaseException] | None,
  196. exc_val: BaseException | None,
  197. exc_tb: TracebackType | None,
  198. ) -> bool:
  199. try:
  200. # trio.Nursery.__exit__ returns bool; .open_nursery has wrong type
  201. return await self._nursery_manager.__aexit__(exc_type, exc_val, exc_tb) # type: ignore[return-value]
  202. except BaseExceptionGroup as exc:
  203. if not exc.split(trio.Cancelled)[1]:
  204. raise trio.Cancelled._create() from exc
  205. raise
  206. finally:
  207. del exc_val, exc_tb
  208. self._active = False
  209. def _check_active(self, coro: Coroutine | None = None) -> None:
  210. if not self._active:
  211. if coro is not None:
  212. coro.close()
  213. raise RuntimeError(
  214. "This task group is not active; no new tasks can be started."
  215. )
  216. def create_task(
  217. self,
  218. coro: Coroutine[Any, Any, T_co],
  219. *,
  220. name: object = None,
  221. context: Context | None = None,
  222. ) -> TaskHandle[T_co]:
  223. if not isinstance(coro, Coroutine):
  224. raise TypeError(f"expected a coroutine, got {coro.__class__.__qualname__}")
  225. self._check_active(coro)
  226. final_name = get_coro_name(coro, name)
  227. handle = TaskHandle(coro, final_name)
  228. if context is not None:
  229. context.run(
  230. partial(self._nursery.start_soon, handle._run_coro, name=handle.name)
  231. )
  232. else:
  233. self._nursery.start_soon(handle._run_coro, name=handle.name)
  234. return handle
  235. async def start(
  236. self,
  237. func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
  238. *args: Unpack[PosArgsT],
  239. name: object = None,
  240. return_handle: Literal[False, True] = False,
  241. ) -> Any:
  242. handle: TaskHandle[T_co]
  243. async def run_coro_with_task_status(
  244. *, task_status: trio.TaskStatus[Any]
  245. ) -> None:
  246. nonlocal handle
  247. wrapper_task_status = _TrioTaskStatus()
  248. coro = call_for_coroutine(func, args, task_status=wrapper_task_status)
  249. if wrapper_task_status.early_start_value is not empty_start_value:
  250. task_status.started(wrapper_task_status.early_start_value)
  251. else:
  252. wrapper_task_status.real_task_status = task_status
  253. handle = TaskHandle(coro, final_name)
  254. await handle._run_coro()
  255. self._check_active()
  256. final_name = get_callable_name(func, name)
  257. start_value = await self._nursery.start(
  258. run_coro_with_task_status, name=final_name
  259. )
  260. if return_handle:
  261. handle._start_value = start_value
  262. return handle
  263. else:
  264. return start_value
  265. #
  266. # Subprocesses
  267. #
  268. @dataclass(eq=False)
  269. class ReceiveStreamWrapper(abc.ByteReceiveStream):
  270. _stream: trio.abc.ReceiveStream
  271. async def receive(self, max_bytes: int | None = None) -> bytes:
  272. if max_bytes is not None and max_bytes < 1:
  273. raise ValueError("max_bytes must be a positive integer")
  274. try:
  275. data = await self._stream.receive_some(max_bytes)
  276. except trio.ClosedResourceError as exc:
  277. raise ClosedResourceError from exc.__cause__
  278. except trio.BrokenResourceError as exc:
  279. raise BrokenResourceError from exc.__cause__
  280. if data:
  281. return bytes(data)
  282. else:
  283. raise EndOfStream
  284. async def aclose(self) -> None:
  285. await self._stream.aclose()
  286. @dataclass(eq=False)
  287. class SendStreamWrapper(abc.ByteSendStream):
  288. _stream: trio.abc.SendStream
  289. async def send(self, item: bytes) -> None:
  290. try:
  291. await self._stream.send_all(item)
  292. except trio.ClosedResourceError as exc:
  293. raise ClosedResourceError from exc.__cause__
  294. except trio.BrokenResourceError as exc:
  295. raise BrokenResourceError from exc.__cause__
  296. async def aclose(self) -> None:
  297. await self._stream.aclose()
  298. @dataclass(eq=False)
  299. class Process(abc.Process):
  300. _process: trio.Process
  301. _stdin: abc.ByteSendStream | None
  302. _stdout: abc.ByteReceiveStream | None
  303. _stderr: abc.ByteReceiveStream | None
  304. async def aclose(self) -> None:
  305. with CancelScope(shield=True):
  306. if self._stdin:
  307. await self._stdin.aclose()
  308. if self._stdout:
  309. await self._stdout.aclose()
  310. if self._stderr:
  311. await self._stderr.aclose()
  312. try:
  313. await self.wait()
  314. except BaseException:
  315. self.kill()
  316. with CancelScope(shield=True):
  317. await self.wait()
  318. raise
  319. async def wait(self) -> int:
  320. return await self._process.wait()
  321. def terminate(self) -> None:
  322. self._process.terminate()
  323. def kill(self) -> None:
  324. self._process.kill()
  325. def send_signal(self, signal: Signals) -> None:
  326. self._process.send_signal(signal)
  327. @property
  328. def pid(self) -> int:
  329. return self._process.pid
  330. @property
  331. def returncode(self) -> int | None:
  332. return self._process.returncode
  333. @property
  334. def stdin(self) -> abc.ByteSendStream | None:
  335. return self._stdin
  336. @property
  337. def stdout(self) -> abc.ByteReceiveStream | None:
  338. return self._stdout
  339. @property
  340. def stderr(self) -> abc.ByteReceiveStream | None:
  341. return self._stderr
  342. class _ProcessPoolShutdownInstrument(trio.abc.Instrument):
  343. def after_run(self) -> None:
  344. super().after_run()
  345. current_default_worker_process_limiter: trio.lowlevel.RunVar = RunVar(
  346. "current_default_worker_process_limiter"
  347. )
  348. async def _shutdown_process_pool(workers: set[abc.Process]) -> None:
  349. try:
  350. await trio.sleep(math.inf)
  351. except trio.Cancelled:
  352. for process in workers:
  353. if process.returncode is None:
  354. process.kill()
  355. with CancelScope(shield=True):
  356. for process in workers:
  357. await process.aclose()
  358. #
  359. # Sockets and networking
  360. #
  361. class _TrioSocketMixin(Generic[T_SockAddr]):
  362. def __init__(self, trio_socket: TrioSocketType) -> None:
  363. self._trio_socket = trio_socket
  364. self._closed = False
  365. def _check_closed(self) -> None:
  366. if self._closed:
  367. raise ClosedResourceError
  368. if self._trio_socket.fileno() < 0:
  369. raise BrokenResourceError
  370. @property
  371. def _raw_socket(self) -> socket.socket:
  372. return self._trio_socket._sock # type: ignore[attr-defined]
  373. async def aclose(self) -> None:
  374. if self._trio_socket.fileno() >= 0:
  375. self._closed = True
  376. self._trio_socket.close()
  377. def _convert_socket_error(self, exc: BaseException) -> NoReturn:
  378. if isinstance(exc, trio.ClosedResourceError):
  379. raise ClosedResourceError from exc
  380. elif self._trio_socket.fileno() < 0 and self._closed:
  381. raise ClosedResourceError from None
  382. elif isinstance(exc, OSError):
  383. raise BrokenResourceError from exc
  384. else:
  385. raise exc
  386. class SocketStream(_TrioSocketMixin, abc.SocketStream):
  387. def __init__(self, trio_socket: TrioSocketType) -> None:
  388. super().__init__(trio_socket)
  389. self._receive_guard = ResourceGuard("reading from")
  390. self._send_guard = ResourceGuard("writing to")
  391. async def receive(self, max_bytes: int = 65536) -> bytes:
  392. if max_bytes < 1:
  393. raise ValueError("max_bytes must be a positive integer")
  394. with self._receive_guard:
  395. try:
  396. data = await self._trio_socket.recv(max_bytes)
  397. except BaseException as exc:
  398. self._convert_socket_error(exc)
  399. if data:
  400. return data
  401. else:
  402. raise EndOfStream
  403. async def send(self, item: bytes) -> None:
  404. with self._send_guard:
  405. view = memoryview(item)
  406. while view:
  407. try:
  408. bytes_sent = await self._trio_socket.send(view)
  409. except BaseException as exc:
  410. self._convert_socket_error(exc)
  411. view = view[bytes_sent:]
  412. async def send_eof(self) -> None:
  413. self._trio_socket.shutdown(socket.SHUT_WR)
  414. class UNIXSocketStream(SocketStream, abc.UNIXSocketStream):
  415. async def receive_fds(self, msglen: int, maxfds: int) -> tuple[bytes, list[int]]:
  416. if not isinstance(msglen, int) or msglen < 0:
  417. raise ValueError("msglen must be a non-negative integer")
  418. if not isinstance(maxfds, int) or maxfds < 1:
  419. raise ValueError("maxfds must be a positive integer")
  420. fds = array.array("i")
  421. await trio.lowlevel.checkpoint()
  422. with self._receive_guard:
  423. while True:
  424. try:
  425. message, ancdata, _flags, _addr = await self._trio_socket.recvmsg(
  426. msglen, socket.CMSG_LEN(maxfds * fds.itemsize)
  427. )
  428. except BaseException as exc:
  429. self._convert_socket_error(exc)
  430. else:
  431. if not message and not ancdata:
  432. raise EndOfStream
  433. break
  434. for cmsg_level, cmsg_type, cmsg_data in ancdata:
  435. if cmsg_level != socket.SOL_SOCKET or cmsg_type != socket.SCM_RIGHTS:
  436. raise RuntimeError(
  437. f"Received unexpected ancillary data; message = {message!r}, "
  438. f"cmsg_level = {cmsg_level}, cmsg_type = {cmsg_type}"
  439. )
  440. fds.frombytes(cmsg_data[: len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
  441. return message, list(fds)
  442. async def send_fds(self, message: bytes, fds: Collection[int | IOBase]) -> None:
  443. if not message:
  444. raise ValueError("message must not be empty")
  445. if not fds:
  446. raise ValueError("fds must not be empty")
  447. filenos: list[int] = []
  448. for fd in fds:
  449. if isinstance(fd, int):
  450. filenos.append(fd)
  451. elif isinstance(fd, IOBase):
  452. filenos.append(fd.fileno())
  453. fdarray = array.array("i", filenos)
  454. await trio.lowlevel.checkpoint()
  455. with self._send_guard:
  456. while True:
  457. try:
  458. await self._trio_socket.sendmsg(
  459. [message],
  460. [
  461. (
  462. socket.SOL_SOCKET,
  463. socket.SCM_RIGHTS,
  464. fdarray,
  465. )
  466. ],
  467. )
  468. break
  469. except BaseException as exc:
  470. self._convert_socket_error(exc)
  471. class TCPSocketListener(_TrioSocketMixin, abc.SocketListener):
  472. def __init__(self, raw_socket: socket.socket):
  473. super().__init__(trio.socket.from_stdlib_socket(raw_socket))
  474. self._accept_guard = ResourceGuard("accepting connections from")
  475. async def accept(self) -> SocketStream:
  476. with self._accept_guard:
  477. try:
  478. trio_socket, _addr = await self._trio_socket.accept()
  479. except BaseException as exc:
  480. self._convert_socket_error(exc)
  481. trio_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  482. return SocketStream(trio_socket)
  483. class UNIXSocketListener(_TrioSocketMixin, abc.SocketListener):
  484. def __init__(self, raw_socket: socket.socket):
  485. super().__init__(trio.socket.from_stdlib_socket(raw_socket))
  486. self._accept_guard = ResourceGuard("accepting connections from")
  487. async def accept(self) -> UNIXSocketStream:
  488. with self._accept_guard:
  489. try:
  490. trio_socket, _addr = await self._trio_socket.accept()
  491. except BaseException as exc:
  492. self._convert_socket_error(exc)
  493. return UNIXSocketStream(trio_socket)
  494. class UDPSocket(_TrioSocketMixin[IPSockAddrType], abc.UDPSocket):
  495. def __init__(self, trio_socket: TrioSocketType) -> None:
  496. super().__init__(trio_socket)
  497. self._receive_guard = ResourceGuard("reading from")
  498. self._send_guard = ResourceGuard("writing to")
  499. async def receive(self) -> tuple[bytes, IPSockAddrType]:
  500. with self._receive_guard:
  501. try:
  502. data, addr = await self._trio_socket.recvfrom(65536)
  503. return data, convert_ipv6_sockaddr(addr)
  504. except BaseException as exc:
  505. self._convert_socket_error(exc)
  506. async def send(self, item: UDPPacketType) -> None:
  507. with self._send_guard:
  508. try:
  509. await self._trio_socket.sendto(*item)
  510. except BaseException as exc:
  511. self._convert_socket_error(exc)
  512. class ConnectedUDPSocket(_TrioSocketMixin[IPSockAddrType], abc.ConnectedUDPSocket):
  513. def __init__(self, trio_socket: TrioSocketType) -> None:
  514. super().__init__(trio_socket)
  515. self._receive_guard = ResourceGuard("reading from")
  516. self._send_guard = ResourceGuard("writing to")
  517. async def receive(self) -> bytes:
  518. with self._receive_guard:
  519. try:
  520. return await self._trio_socket.recv(65536)
  521. except BaseException as exc:
  522. self._convert_socket_error(exc)
  523. async def send(self, item: bytes) -> None:
  524. with self._send_guard:
  525. try:
  526. await self._trio_socket.send(item)
  527. except BaseException as exc:
  528. self._convert_socket_error(exc)
  529. class UNIXDatagramSocket(_TrioSocketMixin[str], abc.UNIXDatagramSocket):
  530. def __init__(self, trio_socket: TrioSocketType) -> None:
  531. super().__init__(trio_socket)
  532. self._receive_guard = ResourceGuard("reading from")
  533. self._send_guard = ResourceGuard("writing to")
  534. async def receive(self) -> UNIXDatagramPacketType:
  535. with self._receive_guard:
  536. try:
  537. data, addr = await self._trio_socket.recvfrom(65536)
  538. return data, addr
  539. except BaseException as exc:
  540. self._convert_socket_error(exc)
  541. async def send(self, item: UNIXDatagramPacketType) -> None:
  542. with self._send_guard:
  543. try:
  544. await self._trio_socket.sendto(*item)
  545. except BaseException as exc:
  546. self._convert_socket_error(exc)
  547. class ConnectedUNIXDatagramSocket(
  548. _TrioSocketMixin[str], abc.ConnectedUNIXDatagramSocket
  549. ):
  550. def __init__(self, trio_socket: TrioSocketType) -> None:
  551. super().__init__(trio_socket)
  552. self._receive_guard = ResourceGuard("reading from")
  553. self._send_guard = ResourceGuard("writing to")
  554. async def receive(self) -> bytes:
  555. with self._receive_guard:
  556. try:
  557. return await self._trio_socket.recv(65536)
  558. except BaseException as exc:
  559. self._convert_socket_error(exc)
  560. async def send(self, item: bytes) -> None:
  561. with self._send_guard:
  562. try:
  563. await self._trio_socket.send(item)
  564. except BaseException as exc:
  565. self._convert_socket_error(exc)
  566. #
  567. # Synchronization
  568. #
  569. class Event(BaseEvent):
  570. __slots__ = ("__original",)
  571. def __new__(cls) -> Self:
  572. return object.__new__(cls)
  573. def __init__(self) -> None:
  574. self.__original = trio.Event()
  575. def is_set(self) -> bool:
  576. return self.__original.is_set()
  577. async def wait(self) -> None:
  578. return await self.__original.wait()
  579. def statistics(self) -> EventStatistics:
  580. orig_statistics = self.__original.statistics()
  581. return EventStatistics(tasks_waiting=orig_statistics.tasks_waiting)
  582. def set(self) -> None:
  583. self.__original.set()
  584. class Lock(BaseLock):
  585. __slots__ = "__original", "_fast_acquire"
  586. def __new__(cls, *, fast_acquire: bool = False) -> Self:
  587. return object.__new__(cls)
  588. def __init__(self, *, fast_acquire: bool = False) -> None:
  589. self._fast_acquire = fast_acquire
  590. self.__original = trio.Lock()
  591. @staticmethod
  592. def _convert_runtime_error_msg(exc: RuntimeError) -> None:
  593. if exc.args == ("attempt to re-acquire an already held Lock",):
  594. exc.args = ("Attempted to acquire an already held Lock",)
  595. async def acquire(self) -> None:
  596. if not self._fast_acquire:
  597. try:
  598. await self.__original.acquire()
  599. except RuntimeError as exc:
  600. self._convert_runtime_error_msg(exc)
  601. raise
  602. return
  603. # This is the "fast path" where we don't let other tasks run
  604. await trio.lowlevel.checkpoint_if_cancelled()
  605. try:
  606. self.__original.acquire_nowait()
  607. except trio.WouldBlock:
  608. await self.__original._lot.park()
  609. except RuntimeError as exc:
  610. self._convert_runtime_error_msg(exc)
  611. raise
  612. def acquire_nowait(self) -> None:
  613. try:
  614. self.__original.acquire_nowait()
  615. except trio.WouldBlock:
  616. raise WouldBlock from None
  617. except RuntimeError as exc:
  618. self._convert_runtime_error_msg(exc)
  619. raise
  620. def locked(self) -> bool:
  621. return self.__original.locked()
  622. def release(self) -> None:
  623. self.__original.release()
  624. def statistics(self) -> LockStatistics:
  625. orig_statistics = self.__original.statistics()
  626. owner = TrioTaskInfo(orig_statistics.owner) if orig_statistics.owner else None
  627. return LockStatistics(
  628. orig_statistics.locked, owner, orig_statistics.tasks_waiting
  629. )
  630. class Semaphore(BaseSemaphore):
  631. __slots__ = ("__original",)
  632. def __new__(
  633. cls,
  634. initial_value: int,
  635. *,
  636. max_value: int | None = None,
  637. fast_acquire: bool = False,
  638. ) -> Self:
  639. return object.__new__(cls)
  640. def __init__(
  641. self,
  642. initial_value: int,
  643. *,
  644. max_value: int | None = None,
  645. fast_acquire: bool = False,
  646. ) -> None:
  647. super().__init__(initial_value, max_value=max_value, fast_acquire=fast_acquire)
  648. self.__original = trio.Semaphore(initial_value, max_value=max_value)
  649. async def acquire(self) -> None:
  650. if not self._fast_acquire:
  651. await self.__original.acquire()
  652. return
  653. # This is the "fast path" where we don't let other tasks run
  654. await trio.lowlevel.checkpoint_if_cancelled()
  655. try:
  656. self.__original.acquire_nowait()
  657. except trio.WouldBlock:
  658. await self.__original._lot.park()
  659. def acquire_nowait(self) -> None:
  660. try:
  661. self.__original.acquire_nowait()
  662. except trio.WouldBlock:
  663. raise WouldBlock from None
  664. @property
  665. def max_value(self) -> int | None:
  666. return self.__original.max_value
  667. @property
  668. def value(self) -> int:
  669. return self.__original.value
  670. def release(self) -> None:
  671. self.__original.release()
  672. def statistics(self) -> SemaphoreStatistics:
  673. orig_statistics = self.__original.statistics()
  674. return SemaphoreStatistics(orig_statistics.tasks_waiting)
  675. class CapacityLimiter(BaseCapacityLimiter):
  676. __slots__ = ("__original",)
  677. def __new__(
  678. cls,
  679. total_tokens: float | None = None,
  680. *,
  681. original: trio.CapacityLimiter | None = None,
  682. ) -> Self:
  683. return object.__new__(cls)
  684. def __init__(
  685. self,
  686. total_tokens: float | None = None,
  687. *,
  688. original: trio.CapacityLimiter | None = None,
  689. ) -> None:
  690. if original is not None:
  691. self.__original = original
  692. else:
  693. assert total_tokens is not None
  694. self.__original = trio.CapacityLimiter(total_tokens)
  695. async def __aenter__(self) -> None:
  696. return await self.__original.__aenter__()
  697. async def __aexit__(
  698. self,
  699. exc_type: type[BaseException] | None,
  700. exc_val: BaseException | None,
  701. exc_tb: TracebackType | None,
  702. ) -> None:
  703. await self.__original.__aexit__(exc_type, exc_val, exc_tb)
  704. @property
  705. def total_tokens(self) -> float:
  706. return self.__original.total_tokens
  707. @total_tokens.setter
  708. def total_tokens(self, value: float) -> None:
  709. self.__original.total_tokens = value
  710. @property
  711. def borrowed_tokens(self) -> int:
  712. return self.__original.borrowed_tokens
  713. @property
  714. def available_tokens(self) -> float:
  715. return self.__original.available_tokens
  716. def acquire_nowait(self) -> None:
  717. try:
  718. self.__original.acquire_nowait()
  719. except trio.WouldBlock:
  720. raise WouldBlock from None
  721. def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
  722. try:
  723. self.__original.acquire_on_behalf_of_nowait(borrower)
  724. except trio.WouldBlock:
  725. raise WouldBlock from None
  726. async def acquire(self) -> None:
  727. await self.__original.acquire()
  728. async def acquire_on_behalf_of(self, borrower: object) -> None:
  729. await self.__original.acquire_on_behalf_of(borrower)
  730. def release(self) -> None:
  731. return self.__original.release()
  732. def release_on_behalf_of(self, borrower: object) -> None:
  733. return self.__original.release_on_behalf_of(borrower)
  734. def statistics(self) -> CapacityLimiterStatistics:
  735. orig = self.__original.statistics()
  736. return CapacityLimiterStatistics(
  737. borrowed_tokens=orig.borrowed_tokens,
  738. total_tokens=orig.total_tokens,
  739. borrowers=tuple(orig.borrowers),
  740. tasks_waiting=orig.tasks_waiting,
  741. )
  742. _capacity_limiter_wrapper: trio.lowlevel.RunVar = RunVar("_capacity_limiter_wrapper")
  743. #
  744. # Signal handling
  745. #
  746. class _SignalReceiver:
  747. _iterator: AsyncIterator[int]
  748. def __init__(self, signals: tuple[Signals, ...]):
  749. self._signals = signals
  750. def __enter__(self) -> Self:
  751. self._cm = trio.open_signal_receiver(*self._signals)
  752. self._iterator = self._cm.__enter__()
  753. return self
  754. def __exit__(
  755. self,
  756. exc_type: type[BaseException] | None,
  757. exc_val: BaseException | None,
  758. exc_tb: TracebackType | None,
  759. ) -> bool | None:
  760. return self._cm.__exit__(exc_type, exc_val, exc_tb)
  761. def __aiter__(self) -> _SignalReceiver:
  762. return self
  763. async def __anext__(self) -> Signals:
  764. signum = await self._iterator.__anext__()
  765. return Signals(signum)
  766. #
  767. # Testing and debugging
  768. #
  769. class TestRunner(abc.TestRunner):
  770. def __init__(self, **options: Any) -> None:
  771. from queue import Queue
  772. self._call_queue: Queue[Callable[[], object]] = Queue()
  773. self._send_stream: (
  774. MemoryObjectSendStream[tuple[Awaitable[Any], list[Outcome]]] | None
  775. ) = None
  776. self._options = options
  777. def __exit__(
  778. self,
  779. exc_type: type[BaseException] | None,
  780. exc_val: BaseException | None,
  781. exc_tb: types.TracebackType | None,
  782. ) -> None:
  783. if self._send_stream:
  784. self._send_stream.close()
  785. while self._send_stream is not None:
  786. self._call_queue.get()()
  787. def is_running(self) -> bool:
  788. return trio.lowlevel.in_trio_task()
  789. async def _run_tests_and_fixtures(self) -> None:
  790. self._send_stream, receive_stream = create_memory_object_stream[
  791. tuple[Awaitable[Any], list[Outcome]]
  792. ](1)
  793. with receive_stream:
  794. async for awaitable, outcome_holder in receive_stream:
  795. try:
  796. retval = await awaitable
  797. except BaseException as exc:
  798. outcome_holder.append(Error(exc))
  799. else:
  800. outcome_holder.append(Value(retval))
  801. def _main_task_finished(self, outcome: object) -> None:
  802. self._send_stream = None
  803. def _call_in_runner_task(
  804. self,
  805. func: Callable[P, Awaitable[T_Retval]],
  806. /,
  807. *args: P.args,
  808. **kwargs: P.kwargs,
  809. ) -> T_Retval:
  810. if self._send_stream is None:
  811. trio.lowlevel.start_guest_run(
  812. self._run_tests_and_fixtures,
  813. run_sync_soon_threadsafe=self._call_queue.put,
  814. done_callback=self._main_task_finished,
  815. **self._options,
  816. )
  817. while self._send_stream is None:
  818. self._call_queue.get()()
  819. outcome_holder: list[Outcome] = []
  820. self._send_stream.send_nowait((func(*args, **kwargs), outcome_holder))
  821. while not outcome_holder:
  822. self._call_queue.get()()
  823. return outcome_holder[0].unwrap()
  824. def run_asyncgen_fixture(
  825. self,
  826. fixture_func: Callable[..., AsyncGenerator[T_Retval, Any]],
  827. kwargs: dict[str, Any],
  828. ) -> Iterable[T_Retval]:
  829. asyncgen = fixture_func(**kwargs)
  830. fixturevalue: T_Retval = self._call_in_runner_task(asyncgen.asend, None)
  831. yield fixturevalue
  832. try:
  833. self._call_in_runner_task(asyncgen.asend, None)
  834. except StopAsyncIteration:
  835. pass
  836. else:
  837. self._call_in_runner_task(asyncgen.aclose)
  838. raise RuntimeError("Async generator fixture did not stop")
  839. def run_fixture(
  840. self,
  841. fixture_func: Callable[..., Coroutine[Any, Any, T_Retval]],
  842. kwargs: dict[str, Any],
  843. ) -> T_Retval:
  844. return self._call_in_runner_task(fixture_func, **kwargs)
  845. def run_test(
  846. self, test_func: Callable[..., Coroutine[Any, Any, Any]], kwargs: dict[str, Any]
  847. ) -> None:
  848. self._call_in_runner_task(test_func, **kwargs)
  849. class TrioTaskInfo(TaskInfo):
  850. def __init__(self, task: trio.lowlevel.Task):
  851. parent_id = None
  852. if task.parent_nursery and task.parent_nursery.parent_task:
  853. parent_id = id(task.parent_nursery.parent_task)
  854. super().__init__(id(task), parent_id, task.name, task.coro)
  855. self._task = weakref.proxy(task)
  856. def has_pending_cancellation(self) -> bool:
  857. try:
  858. return self._task._cancel_status.effectively_cancelled
  859. except ReferenceError:
  860. # If the task is no longer around, it surely doesn't have a cancellation
  861. # pending
  862. return False
  863. class TrioBackend(AsyncBackend):
  864. @classmethod
  865. def run(
  866. cls,
  867. func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval]],
  868. args: tuple[Unpack[PosArgsT]],
  869. kwargs: dict[str, Any],
  870. options: dict[str, Any],
  871. ) -> T_Retval:
  872. assert not kwargs, "unreachable, and not supported by Trio"
  873. return trio.run(ensure_returns_coro(func), *args, **options)
  874. @classmethod
  875. def current_token(cls) -> object:
  876. return trio.lowlevel.current_trio_token()
  877. @classmethod
  878. def current_time(cls) -> float:
  879. return trio.current_time()
  880. @classmethod
  881. def cancelled_exception_class(cls) -> type[BaseException]:
  882. return trio.Cancelled
  883. @classmethod
  884. async def checkpoint(cls) -> None:
  885. await trio.lowlevel.checkpoint()
  886. @classmethod
  887. async def checkpoint_if_cancelled(cls) -> None:
  888. await trio.lowlevel.checkpoint_if_cancelled()
  889. @classmethod
  890. async def cancel_shielded_checkpoint(cls) -> None:
  891. await trio.lowlevel.cancel_shielded_checkpoint()
  892. @classmethod
  893. async def sleep(cls, delay: float) -> None:
  894. await trio.sleep(delay)
  895. @classmethod
  896. def create_cancel_scope(
  897. cls, *, deadline: float = math.inf, shield: bool = False
  898. ) -> BaseCancelScope:
  899. return CancelScope(deadline=deadline, shield=shield)
  900. @classmethod
  901. def current_effective_deadline(cls) -> float:
  902. return trio.current_effective_deadline()
  903. @classmethod
  904. def create_task_group(cls) -> abc.TaskGroup:
  905. return TaskGroup()
  906. @classmethod
  907. def create_event(cls) -> BaseEvent:
  908. return Event()
  909. @classmethod
  910. def create_lock(cls, *, fast_acquire: bool) -> Lock:
  911. return Lock(fast_acquire=fast_acquire)
  912. @classmethod
  913. def create_semaphore(
  914. cls,
  915. initial_value: int,
  916. *,
  917. max_value: int | None = None,
  918. fast_acquire: bool = False,
  919. ) -> BaseSemaphore:
  920. return Semaphore(initial_value, max_value=max_value, fast_acquire=fast_acquire)
  921. @classmethod
  922. def create_capacity_limiter(cls, total_tokens: float) -> CapacityLimiter:
  923. return CapacityLimiter(total_tokens)
  924. @classmethod
  925. async def run_sync_in_worker_thread(
  926. cls,
  927. func: Callable[[Unpack[PosArgsT]], T_Retval],
  928. args: tuple[Unpack[PosArgsT]],
  929. abandon_on_cancel: bool = False,
  930. limiter: BaseCapacityLimiter | None = None,
  931. ) -> T_Retval:
  932. def wrapper() -> T_Retval:
  933. with claim_worker_thread(TrioBackend, token):
  934. return func(*args)
  935. token = TrioBackend.current_token()
  936. return await run_sync(
  937. wrapper,
  938. abandon_on_cancel=abandon_on_cancel,
  939. limiter=cast(trio.CapacityLimiter, limiter),
  940. )
  941. @classmethod
  942. def check_cancelled(cls) -> None:
  943. trio.from_thread.check_cancelled()
  944. @classmethod
  945. def run_async_from_thread(
  946. cls,
  947. func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
  948. args: tuple[Unpack[PosArgsT]],
  949. token: object,
  950. ) -> T_co:
  951. trio_token = cast("trio.lowlevel.TrioToken | None", token)
  952. try:
  953. return trio.from_thread.run(func, *args, trio_token=trio_token)
  954. except trio.RunFinishedError:
  955. raise RunFinishedError from None
  956. @classmethod
  957. def run_sync_from_thread(
  958. cls,
  959. func: Callable[[Unpack[PosArgsT]], T_Retval],
  960. args: tuple[Unpack[PosArgsT]],
  961. token: object,
  962. ) -> T_Retval:
  963. trio_token = cast("trio.lowlevel.TrioToken | None", token)
  964. try:
  965. return trio.from_thread.run_sync(func, *args, trio_token=trio_token)
  966. except trio.RunFinishedError:
  967. raise RunFinishedError from None
  968. @classmethod
  969. async def open_process(
  970. cls,
  971. command: StrOrBytesPath | Sequence[StrOrBytesPath],
  972. *,
  973. stdin: int | IO[Any] | None,
  974. stdout: int | IO[Any] | None,
  975. stderr: int | IO[Any] | None,
  976. **kwargs: Any,
  977. ) -> Process:
  978. def convert_item(item: StrOrBytesPath) -> str:
  979. str_or_bytes = os.fspath(item)
  980. if isinstance(str_or_bytes, str):
  981. return str_or_bytes
  982. else:
  983. return os.fsdecode(str_or_bytes)
  984. if isinstance(command, (str, bytes, PathLike)):
  985. process = await trio.lowlevel.open_process(
  986. convert_item(command),
  987. stdin=stdin,
  988. stdout=stdout,
  989. stderr=stderr,
  990. shell=True,
  991. **kwargs,
  992. )
  993. else:
  994. process = await trio.lowlevel.open_process(
  995. [convert_item(item) for item in command],
  996. stdin=stdin,
  997. stdout=stdout,
  998. stderr=stderr,
  999. shell=False,
  1000. **kwargs,
  1001. )
  1002. stdin_stream = SendStreamWrapper(process.stdin) if process.stdin else None
  1003. stdout_stream = ReceiveStreamWrapper(process.stdout) if process.stdout else None
  1004. stderr_stream = ReceiveStreamWrapper(process.stderr) if process.stderr else None
  1005. return Process(process, stdin_stream, stdout_stream, stderr_stream)
  1006. @classmethod
  1007. def setup_process_pool_exit_at_shutdown(cls, workers: set[abc.Process]) -> None:
  1008. trio.lowlevel.spawn_system_task(_shutdown_process_pool, workers)
  1009. @classmethod
  1010. async def connect_tcp(
  1011. cls, host: str, port: int, local_address: IPSockAddrType | None = None
  1012. ) -> SocketStream:
  1013. family = socket.AF_INET6 if ":" in host else socket.AF_INET
  1014. trio_socket = trio.socket.socket(family)
  1015. trio_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  1016. if local_address:
  1017. await trio_socket.bind(local_address)
  1018. try:
  1019. await trio_socket.connect((host, port))
  1020. except BaseException:
  1021. trio_socket.close()
  1022. raise
  1023. return SocketStream(trio_socket)
  1024. @classmethod
  1025. async def connect_unix(cls, path: str | bytes) -> abc.UNIXSocketStream:
  1026. trio_socket = trio.socket.socket(socket.AF_UNIX)
  1027. try:
  1028. await trio_socket.connect(path)
  1029. except BaseException:
  1030. trio_socket.close()
  1031. raise
  1032. return UNIXSocketStream(trio_socket)
  1033. @classmethod
  1034. def create_tcp_listener(cls, sock: socket.socket) -> abc.SocketListener:
  1035. return TCPSocketListener(sock)
  1036. @classmethod
  1037. def create_unix_listener(cls, sock: socket.socket) -> abc.SocketListener:
  1038. return UNIXSocketListener(sock)
  1039. @classmethod
  1040. async def create_udp_socket(
  1041. cls,
  1042. family: socket.AddressFamily,
  1043. local_address: IPSockAddrType | None,
  1044. remote_address: IPSockAddrType | None,
  1045. reuse_port: bool,
  1046. ) -> UDPSocket | ConnectedUDPSocket:
  1047. trio_socket = trio.socket.socket(family=family, type=socket.SOCK_DGRAM)
  1048. if reuse_port:
  1049. trio_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
  1050. if local_address:
  1051. await trio_socket.bind(local_address)
  1052. if remote_address:
  1053. await trio_socket.connect(remote_address)
  1054. return ConnectedUDPSocket(trio_socket)
  1055. else:
  1056. return UDPSocket(trio_socket)
  1057. @classmethod
  1058. @overload
  1059. async def create_unix_datagram_socket(
  1060. cls, raw_socket: socket.socket, remote_path: None
  1061. ) -> abc.UNIXDatagramSocket: ...
  1062. @classmethod
  1063. @overload
  1064. async def create_unix_datagram_socket(
  1065. cls, raw_socket: socket.socket, remote_path: str | bytes
  1066. ) -> abc.ConnectedUNIXDatagramSocket: ...
  1067. @classmethod
  1068. async def create_unix_datagram_socket(
  1069. cls, raw_socket: socket.socket, remote_path: str | bytes | None
  1070. ) -> abc.UNIXDatagramSocket | abc.ConnectedUNIXDatagramSocket:
  1071. trio_socket = trio.socket.from_stdlib_socket(raw_socket)
  1072. if remote_path:
  1073. await trio_socket.connect(remote_path)
  1074. return ConnectedUNIXDatagramSocket(trio_socket)
  1075. else:
  1076. return UNIXDatagramSocket(trio_socket)
  1077. @classmethod
  1078. async def getaddrinfo(
  1079. cls,
  1080. host: bytes | str | None,
  1081. port: str | int | None,
  1082. *,
  1083. family: int | AddressFamily = 0,
  1084. type: int | SocketKind = 0,
  1085. proto: int = 0,
  1086. flags: int = 0,
  1087. ) -> Sequence[
  1088. tuple[
  1089. AddressFamily,
  1090. SocketKind,
  1091. int,
  1092. str,
  1093. tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes],
  1094. ]
  1095. ]:
  1096. return await trio.socket.getaddrinfo(host, port, family, type, proto, flags)
  1097. @classmethod
  1098. async def getnameinfo(
  1099. cls, sockaddr: IPSockAddrType, flags: int = 0
  1100. ) -> tuple[str, str]:
  1101. return await trio.socket.getnameinfo(sockaddr, flags)
  1102. @classmethod
  1103. async def wait_readable(cls, obj: FileDescriptorLike) -> None:
  1104. try:
  1105. await wait_readable(obj)
  1106. except trio.ClosedResourceError as exc:
  1107. raise ClosedResourceError().with_traceback(exc.__traceback__) from None
  1108. except trio.BusyResourceError:
  1109. raise BusyResourceError("reading from") from None
  1110. @classmethod
  1111. async def wait_writable(cls, obj: FileDescriptorLike) -> None:
  1112. try:
  1113. await wait_writable(obj)
  1114. except trio.ClosedResourceError as exc:
  1115. raise ClosedResourceError().with_traceback(exc.__traceback__) from None
  1116. except trio.BusyResourceError:
  1117. raise BusyResourceError("writing to") from None
  1118. @classmethod
  1119. def notify_closing(cls, obj: FileDescriptorLike) -> None:
  1120. notify_closing(obj)
  1121. @classmethod
  1122. async def wrap_listener_socket(cls, sock: socket.socket) -> abc.SocketListener:
  1123. if hasattr(socket, "AF_UNIX") and sock.family == socket.AF_UNIX:
  1124. return UNIXSocketListener(sock)
  1125. return TCPSocketListener(sock)
  1126. @classmethod
  1127. async def wrap_stream_socket(cls, sock: socket.socket) -> SocketStream:
  1128. trio_sock = trio.socket.from_stdlib_socket(sock)
  1129. return SocketStream(trio_sock)
  1130. @classmethod
  1131. async def wrap_unix_stream_socket(cls, sock: socket.socket) -> UNIXSocketStream:
  1132. trio_sock = trio.socket.from_stdlib_socket(sock)
  1133. return UNIXSocketStream(trio_sock)
  1134. @classmethod
  1135. async def wrap_udp_socket(cls, sock: socket.socket) -> UDPSocket:
  1136. trio_sock = trio.socket.from_stdlib_socket(sock)
  1137. return UDPSocket(trio_sock)
  1138. @classmethod
  1139. async def wrap_connected_udp_socket(cls, sock: socket.socket) -> ConnectedUDPSocket:
  1140. trio_sock = trio.socket.from_stdlib_socket(sock)
  1141. return ConnectedUDPSocket(trio_sock)
  1142. @classmethod
  1143. async def wrap_unix_datagram_socket(cls, sock: socket.socket) -> UNIXDatagramSocket:
  1144. trio_sock = trio.socket.from_stdlib_socket(sock)
  1145. return UNIXDatagramSocket(trio_sock)
  1146. @classmethod
  1147. async def wrap_connected_unix_datagram_socket(
  1148. cls, sock: socket.socket
  1149. ) -> ConnectedUNIXDatagramSocket:
  1150. trio_sock = trio.socket.from_stdlib_socket(sock)
  1151. return ConnectedUNIXDatagramSocket(trio_sock)
  1152. @classmethod
  1153. def current_default_thread_limiter(cls) -> CapacityLimiter:
  1154. try:
  1155. return _capacity_limiter_wrapper.get()
  1156. except LookupError:
  1157. limiter = CapacityLimiter(
  1158. original=trio.to_thread.current_default_thread_limiter()
  1159. )
  1160. _capacity_limiter_wrapper.set(limiter)
  1161. return limiter
  1162. @classmethod
  1163. def open_signal_receiver(
  1164. cls, *signals: Signals
  1165. ) -> AbstractContextManager[AsyncIterator[Signals]]:
  1166. return _SignalReceiver(signals)
  1167. @classmethod
  1168. def get_current_task(cls) -> TaskInfo:
  1169. task = current_task()
  1170. return TrioTaskInfo(task)
  1171. @classmethod
  1172. def get_running_tasks(cls) -> Sequence[TaskInfo]:
  1173. root_task = current_root_task()
  1174. assert root_task
  1175. task_infos = [TrioTaskInfo(root_task)]
  1176. nurseries = root_task.child_nurseries
  1177. while nurseries:
  1178. new_nurseries: list[trio.Nursery] = []
  1179. for nursery in nurseries:
  1180. for task in nursery.child_tasks:
  1181. task_infos.append(TrioTaskInfo(task))
  1182. new_nurseries.extend(task.child_nurseries)
  1183. nurseries = new_nurseries
  1184. return task_infos
  1185. @classmethod
  1186. async def wait_all_tasks_blocked(cls) -> None:
  1187. from trio.testing import wait_all_tasks_blocked
  1188. await wait_all_tasks_blocked()
  1189. @classmethod
  1190. def create_test_runner(cls, options: dict[str, Any]) -> TestRunner:
  1191. return TestRunner(**options)
  1192. backend_class = TrioBackend