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

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(ObjectStream[T_Item], Generic[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 def send_nowait(self, item: T_Item) -> None: 

89 try: 

90 send_nowait = self.send_stream.send_nowait # type: ignore[attr-defined] 

91 except AttributeError as exc: 

92 raise NotImplementedError( 

93 f"'send_nowait' method not implemented in {type(self.send_stream)}" 

94 ) from exc 

95 

96 send_nowait(item) 

97 

98 async def send_eof(self) -> None: 

99 await self.send_stream.aclose() 

100 

101 async def aclose(self) -> None: 

102 await self.send_stream.aclose() 

103 await self.receive_stream.aclose() 

104 

105 @property 

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

107 return { 

108 **self.send_stream.extra_attributes, 

109 **self.receive_stream.extra_attributes, 

110 } 

111 

112 

113@dataclass(eq=False) 

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

115 """ 

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

117 

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

119 moved into this one. 

120 

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

122 overriding any conflicting attributes from the previous one. 

123 

124 :param listeners: listeners to serve 

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

126 """ 

127 

128 listeners: Sequence[Listener[T_Stream]] 

129 

130 def __post_init__(self) -> None: 

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

132 for listener in self.listeners: 

133 if isinstance(listener, MultiListener): 

134 listeners.extend(listener.listeners) 

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

136 else: 

137 listeners.append(listener) 

138 

139 self.listeners = listeners 

140 

141 async def serve( 

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

143 ) -> None: 

144 from .. import create_task_group 

145 

146 async with create_task_group() as tg: 

147 for listener in self.listeners: 

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

149 

150 async def aclose(self) -> None: 

151 for listener in self.listeners: 

152 await listener.aclose() 

153 

154 @property 

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

156 attributes: dict = {} 

157 for listener in self.listeners: 

158 attributes.update(listener.extra_attributes) 

159 

160 return attributes