Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/streams/common.py: 42%
26 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
1import asyncio
2import logging
3from contextlib import asynccontextmanager, suppress
5from anyio import (
6 BrokenResourceError,
7 ClosedResourceError,
8 create_memory_object_stream,
9 create_task_group,
10)
11from anyio.abc import AnyByteStream
12from anyio.streams.stapled import StapledObjectStream
14logger = logging.getLogger(__name__)
17async def copy_stream(dst: AnyByteStream, src: AnyByteStream):
18 with suppress(BrokenResourceError, ClosedResourceError, asyncio.exceptions.InvalidStateError):
19 async for v in src:
20 await dst.send(v)
21 with suppress(AttributeError):
22 await dst.send_eof()
25@asynccontextmanager
26async def forward_stream(a, b):
27 async with a, b:
28 async with create_task_group() as tg:
29 tg.start_soon(copy_stream, a, b)
30 tg.start_soon(copy_stream, b, a)
31 yield
34def create_memory_stream():
35 a_tx, a_rx = create_memory_object_stream[bytes](32)
36 b_tx, b_rx = create_memory_object_stream[bytes](32)
37 a = StapledObjectStream(a_tx, b_rx)
38 b = StapledObjectStream(b_tx, a_rx)
39 return a, b