Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/streams/metadata.py: 68%

25 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-05-06 10:20 +0200

1from contextlib import suppress 

2from dataclasses import dataclass 

3from typing import Any, Callable, Mapping 

4 

5from anyio import TypedAttributeLookupError, TypedAttributeSet, typed_attribute 

6from anyio.abc import AnyByteStream, ObjectStream 

7 

8 

9class MetadataStreamAttributes(TypedAttributeSet): 

10 # https://grpc.io/docs/guides/metadata/ 

11 metadata: dict[str, str] = typed_attribute() 

12 

13 

14@dataclass(frozen=True, kw_only=True, slots=True) 

15class MetadataStream(ObjectStream[bytes]): 

16 stream: AnyByteStream 

17 metadata: dict[str, str] 

18 

19 async def send(self, item: bytes): 

20 await self.stream.send(item) 

21 

22 async def receive(self) -> bytes: 

23 return await self.stream.receive() 

24 

25 async def send_eof(self): 

26 await self.stream.send_eof() 

27 

28 async def aclose(self): 

29 await self.stream.aclose() 

30 

31 @property 

32 def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]: 

33 metadata = {} 

34 with suppress(TypedAttributeLookupError): 

35 metadata = self.stream.extra(MetadataStreamAttributes.metadata) 

36 return self.stream.extra_attributes | {MetadataStreamAttributes.metadata: lambda: metadata | self.metadata}