Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/websockets/streams.py: 75%

1from __future__ import annotations 

2 

3from collections.abc import Generator 

4 

5 

6class StreamReader: 

7 """ 

8 Generator-based stream reader. 

9 

10 This class doesn't support concurrent calls to :meth:`read_line`, 

11 :meth:`read_exact`, or :meth:`read_to_eof`. Make sure calls are 

12 serialized. 

13 

14 """ 

15 

16 def __init__(self) -> None: 

17 self.buffer = bytearray() 

18 self.eof = False 

19 

20 def read_line( 

21 self, 

22 m: int, 

23 too_long_exc_type: type[Exception] = RuntimeError, 

24 ) -> Generator[None, None, bytearray]: 

25 """ 

26 Read a LF-terminated line from the stream. 

27 

28 This is a generator-based coroutine. 

29 

30 The return value includes the LF character. 

31 

32 Args: 

33 m: Maximum number bytes to read; this is a security limit. 

34 too_long_exc_type: exception to raise if the line ends in more 

35 than ``m`` bytes; defaults to :exc:`RuntimeError`. 

36 

37 Raises: 

38 EOFError: If the stream ends without a LF. 

39 RuntimeError: If the line ends in more than ``m`` bytes. 

40 

41 """ 

42 n = 0 # number of bytes to read 

43 p = 0 # number of bytes without a newline 

44 while True: 

45 n = self.buffer.find(b"\n", p) + 1 

46 if n > 0: 

47 break 

48 p = len(self.buffer) 

49 if p > m: 

50 raise too_long_exc_type( 

51 f"read {p} bytes, expected no more than {m} bytes" 

52 ) 

53 if self.eof: 

54 raise EOFError(f"stream ends after {p} bytes, before end of line") 

55 yield 

56 if n > m: 

57 raise too_long_exc_type(f"read {n} bytes, expected no more than {m} bytes") 

58 r = self.buffer[:n] 

59 del self.buffer[:n] 

60 return r 

61 

62 def read_exact(self, n: int) -> Generator[None, None, bytearray]: 

63 """ 

64 Read a given number of bytes from the stream. 

65 

66 This is a generator-based coroutine. 

67 

68 Args: 

69 n: How many bytes to read. 

70 

71 Raises: 

72 EOFError: If the stream ends in less than ``n`` bytes. 

73 

74 """ 

75 assert n >= 0 

76 while len(self.buffer) < n: 

77 if self.eof: 

78 p = len(self.buffer) 

79 raise EOFError(f"stream ends after {p} bytes, expected {n} bytes") 

80 yield 

81 r = self.buffer[:n] 

82 del self.buffer[:n] 

83 return r 

84 

85 def read_to_eof( 

86 self, 

87 m: int, 

88 too_long_exc_type: type[Exception] = RuntimeError, 

89 ) -> Generator[None, None, bytearray]: 

90 """ 

91 Read all bytes from the stream. 

92 

93 This is a generator-based coroutine. 

94 

95 Args: 

96 m: Maximum number bytes to read; this is a security limit. 

97 too_long_exc_type: exception to raise if the stream ends in more 

98 than ``m`` bytes; defaults to :exc:`RuntimeError`. 

99 

100 Raises: 

101 RuntimeError: If the stream ends in more than ``m`` bytes. 

102 

103 """ 

104 while not self.eof: 

105 p = len(self.buffer) 

106 if p > m: 

107 raise too_long_exc_type( 

108 f"read {p} bytes, expected no more than {m} bytes" 

109 ) 

110 yield 

111 r = self.buffer[:] 

112 del self.buffer[:] 

113 return r 

114 

115 def at_eof(self) -> Generator[None, None, bool]: 

116 """ 

117 Tell whether the stream has ended and all data was read. 

118 

119 This is a generator-based coroutine. 

120 

121 """ 

122 while True: 

123 if self.buffer: 

124 return False 

125 if self.eof: 

126 return True 

127 # When all data was read but the stream hasn't ended, we can't 

128 # tell if until either feed_data() or feed_eof() is called. 

129 yield 

130 

131 def feed_data(self, data: bytes | bytearray) -> None: 

132 """ 

133 Write data to the stream. 

134 

135 :meth:`feed_data` cannot be called after :meth:`feed_eof`. 

136 

137 Args: 

138 data: Data to write. 

139 

140 Raises: 

141 EOFError: If the stream has ended. 

142 

143 """ 

144 if self.eof: 

145 raise EOFError("stream ended") 

146 self.buffer += data 

147 

148 def feed_eof(self) -> None: 

149 """ 

150 End the stream. 

151 

152 :meth:`feed_eof` cannot be called more than once. 

153 

154 Raises: 

155 EOFError: If the stream has ended. 

156 

157 """ 

158 if self.eof: 

159 raise EOFError("stream ended") 

160 self.eof = True 

161 

162 def discard(self) -> None: 

163 """ 

164 Discard all buffered data, but don't end the stream. 

165 

166 """ 

167 del self.buffer[:]