Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/streams/stapled.py: 54%

1from __future__ import annotations 

2 

3__all__ = ( 

4 "MultiListener", 

5 "StapledByteStream", 

6 "StapledObjectStream", 

7) 

8 

9from collections.abc import Callable, Mapping, Sequence 

10from dataclasses import dataclass 

11from typing import Any, Generic, TypeVar 

12 

13from ..abc import ( 

14 ByteReceiveStream, 

15 ByteSendStream, 

16 ByteStream, 

17 Listener, 

18 ObjectReceiveStream, 

19 ObjectSendStream, 

20 ObjectStream, 

21 TaskGroup, 

22) 

23 

24T_Item = TypeVar("T_Item") 

25T_Stream = TypeVar("T_Stream") 

26 

27 

28@dataclass(eq=False) 

29class StapledByteStream(ByteStream): 

30 """ 

31 Combines two byte streams into a single, bidirectional byte stream. 

32 

33 Extra attributes will be provided from both streams, with the receive stream 

34 providing the values in case of a conflict. 

35 

36 :param ByteSendStream send_stream: the sending byte stream 

37 :param ByteReceiveStream receive_stream: the receiving byte stream 

38 """ 

39 

40 send_stream: ByteSendStream 

41 receive_stream: ByteReceiveStream 

42 

43 async def receive(self, max_bytes: int = 65536) -> bytes: 

44 if max_bytes < 1: 

45 raise ValueError("max_bytes must be a positive integer") 

46 

47 return await self.receive_stream.receive(max_bytes) 

48 

49 async def send(self, item: bytes) -> None: 

50 await self.send_stream.send(item) 

51 

52 async def send_eof(self) -> None: 

53 await self.send_stream.aclose() 

54 

55 async def aclose(self) -> None: 

56 await self.send_stream.aclose() 

57 await self.receive_stream.aclose() 

58 

59 @property 

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

61 return { 

62 **self.send_stream.extra_attributes, 

63 **self.receive_stream.extra_attributes, 

64 } 

65 

66 

67@dataclass(eq=False) 

68class StapledObjectStream(Generic[T_Item], ObjectStream[T_Item]): 

69 """ 

70 Combines two object streams into a single, bidirectional object stream. 

71 

72 Extra attributes will be provided from both streams, with the receive stream 

73 providing the values in case of a conflict. 

74 

75 :param ObjectSendStream send_stream: the sending object stream 

76 :param ObjectReceiveStream receive_stream: the receiving object stream 

77 """ 

78 

79 send_stream: ObjectSendStream[T_Item] 

80 receive_stream: ObjectReceiveStream[T_Item] 

81 

82 async def receive(self) -> T_Item: 

83 return await self.receive_stream.receive() 

84 

85 async def send(self, item: T_Item) -> None: 

86 await self.send_stream.send(item) 

87 

88 async def send_eof(self) -> None: 

89 await self.send_stream.aclose() 

90 

91 async def aclose(self) -> None: 

92 await self.send_stream.aclose() 

93 await self.receive_stream.aclose() 

94 

95 @property 

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

97 return { 

98 **self.send_stream.extra_attributes, 

99 **self.receive_stream.extra_attributes, 

100 } 

101 

102 

103@dataclass(eq=False) 

104class MultiListener(Generic[T_Stream], Listener[T_Stream]): 

105 """ 

106 Combines multiple listeners into one, serving connections from all of them at once. 

107 

108 Any MultiListeners in the given collection of listeners will have their listeners 

109 moved into this one. 

110 

111 Extra attributes are provided from each listener, with each successive listener 

112 overriding any conflicting attributes from the previous one. 

113 

114 :param listeners: listeners to serve 

115 :type listeners: Sequence[Listener[T_Stream]] 

116 """ 

117 

118 listeners: Sequence[Listener[T_Stream]] 

119 

120 def __post_init__(self) -> None: 

121 listeners: list[Listener[T_Stream]] = [] 

122 for listener in self.listeners: 

123 if isinstance(listener, MultiListener): 

124 listeners.extend(listener.listeners) 

125 del listener.listeners[:] # type: ignore[attr-defined] 

126 else: 

127 listeners.append(listener) 

128 

129 self.listeners = listeners 

130 

131 async def serve( 

132 self, handler: Callable[[T_Stream], Any], task_group: TaskGroup | None = None 

133 ) -> None: 

134 from .. import create_task_group 

135 

136 async with create_task_group() as tg: 

137 for listener in self.listeners: 

138 tg.start_soon(listener.serve, handler, task_group) 

139 

140 async def aclose(self) -> None: 

141 for listener in self.listeners: 

142 await listener.aclose() 

143 

144 @property 

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

146 attributes: dict = {} 

147 for listener in self.listeners: 

148 attributes.update(listener.extra_attributes) 

149 

150 return attributes