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

1import asyncio 

2import logging 

3from contextlib import asynccontextmanager, suppress 

4 

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 

13 

14logger = logging.getLogger(__name__) 

15 

16 

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() 

23 

24 

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 

32 

33 

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