Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/client_proto.py: 21%

1import asyncio 

2from contextlib import suppress 

3from typing import Callable, Protocol 

4 

5from ._websocket.reader import WebSocketDataQueue 

6from .base_protocol import BaseProtocol 

7from .client_exceptions import ( 

8 ClientConnectionError, 

9 ClientOSError, 

10 ClientPayloadError, 

11 ServerDisconnectedError, 

12 SocketTimeoutError, 

13) 

14from .helpers import ( 

15 _EXC_SENTINEL, 

16 DEFAULT_CHUNK_SIZE, 

17 EMPTY_BODY_STATUS_CODES, 

18 BaseTimerContext, 

19 ErrorableProtocol, 

20 set_exception, 

21 set_result, 

22) 

23from .http import HttpResponseParser, RawResponseMessage, WebSocketReader 

24from .http_exceptions import HttpProcessingError 

25from .streams import EMPTY_PAYLOAD, DataQueue, StreamReader 

26 

27 

28class _Payload(ErrorableProtocol, Protocol): 

29 def is_eof(self) -> bool: ... 

30 

31 

32class ResponseHandler(BaseProtocol, DataQueue[tuple[RawResponseMessage, StreamReader]]): 

33 """Helper class to adapt between Protocol and StreamReader.""" 

34 

35 def __init__(self, loop: asyncio.AbstractEventLoop) -> None: 

36 BaseProtocol.__init__(self, loop=loop, parser=None) 

37 DataQueue.__init__(self, loop) 

38 

39 self._should_close = False 

40 

41 self._payload: _Payload | None = None 

42 self._skip_payload = False 

43 self._payload_parser: WebSocketReader | None = None 

44 self._data_received_cb: Callable[[], None] | None = None 

45 

46 self._timer = None 

47 self._tail = b"" 

48 

49 self._read_timeout: float | None = None 

50 self._read_timeout_handle: asyncio.TimerHandle | None = None 

51 

52 self._timeout_ceil_threshold: float | None = 5 

53 

54 self._closed: None | asyncio.Future[None] = None 

55 self._connection_lost_called = False 

56 

57 @property 

58 def closed(self) -> None | asyncio.Future[None]: 

59 """Future that is set when the connection is closed. 

60 

61 This property returns a Future that will be completed when the connection 

62 is closed. The Future is created lazily on first access to avoid creating 

63 futures that will never be awaited. 

64 

65 Returns: 

66 - A Future[None] if the connection is still open or was closed after 

67 this property was accessed 

68 - None if connection_lost() was already called before this property 

69 was ever accessed (indicating no one is waiting for the closure) 

70 """ 

71 if self._closed is None and not self._connection_lost_called: 

72 self._closed = self._loop.create_future() 

73 return self._closed 

74 

75 @property 

76 def upgraded(self) -> bool: 

77 return self._upgraded 

78 

79 @property 

80 def should_close(self) -> bool: 

81 return bool( 

82 self._should_close 

83 or (self._payload is not None and not self._payload.is_eof()) 

84 or self._upgraded 

85 or self._exception is not None 

86 or self._payload_parser is not None 

87 or self._buffer 

88 or self._tail 

89 ) 

90 

91 def force_close(self) -> None: 

92 self._should_close = True 

93 

94 def close(self) -> None: 

95 self._exception = None # Break cyclic references 

96 transport = self.transport 

97 if transport is not None: 

98 transport.close() 

99 self.transport = None 

100 self._payload = None 

101 self._drop_timeout() 

102 

103 def abort(self) -> None: 

104 self._exception = None # Break cyclic references 

105 transport = self.transport 

106 if transport is not None: 

107 transport.abort() 

108 self.transport = None 

109 self._payload = None 

110 self._drop_timeout() 

111 

112 def is_connected(self) -> bool: 

113 return self.transport is not None and not self.transport.is_closing() 

114 

115 def connection_lost(self, exc: BaseException | None) -> None: 

116 self._connection_lost_called = True 

117 self._drop_timeout() 

118 

119 original_connection_error = exc 

120 reraised_exc = original_connection_error 

121 

122 connection_closed_cleanly = original_connection_error is None 

123 

124 if self._closed is not None: 

125 # If someone is waiting for the closed future, 

126 # we should set it to None or an exception. If 

127 # self._closed is None, it means that 

128 # connection_lost() was called already 

129 # or nobody is waiting for it. 

130 if connection_closed_cleanly: 

131 set_result(self._closed, None) 

132 else: 

133 assert original_connection_error is not None 

134 set_exception( 

135 self._closed, 

136 ClientConnectionError( 

137 f"Connection lost: {original_connection_error !s}", 

138 ), 

139 original_connection_error, 

140 ) 

141 

142 if self._payload_parser is not None: 

143 with suppress(Exception): # FIXME: log this somehow? 

144 self._payload_parser.feed_eof() 

145 

146 uncompleted = None 

147 if self._parser is not None: 

148 try: 

149 uncompleted = self._parser.feed_eof() 

150 except Exception as underlying_exc: 

151 if self._payload is not None: 

152 client_payload_exc_msg = ( 

153 f"Response payload is not completed: {underlying_exc !r}" 

154 ) 

155 if not connection_closed_cleanly: 

156 client_payload_exc_msg = ( 

157 f"{client_payload_exc_msg !s}. " 

158 f"{original_connection_error !r}" 

159 ) 

160 set_exception( 

161 self._payload, 

162 ClientPayloadError(client_payload_exc_msg), 

163 underlying_exc, 

164 ) 

165 

166 if not self.is_eof(): 

167 if isinstance(original_connection_error, OSError): 

168 reraised_exc = ClientOSError(*original_connection_error.args) 

169 if connection_closed_cleanly: 

170 reraised_exc = ServerDisconnectedError(uncompleted) 

171 # assigns self._should_close to True as side effect, 

172 # we do it anyway below 

173 underlying_non_eof_exc = ( 

174 _EXC_SENTINEL 

175 if connection_closed_cleanly 

176 else original_connection_error 

177 ) 

178 assert underlying_non_eof_exc is not None 

179 assert reraised_exc is not None 

180 self.set_exception(reraised_exc, underlying_non_eof_exc) 

181 

182 self._should_close = True 

183 self._parser = None 

184 self._payload = None 

185 self._payload_parser = None 

186 self._reading_paused = False 

187 

188 super().connection_lost(reraised_exc) 

189 

190 def eof_received(self) -> None: 

191 # should call parser.feed_eof() most likely 

192 self._drop_timeout() 

193 

194 def pause_reading(self) -> None: 

195 super().pause_reading() 

196 self._drop_timeout() 

197 

198 def resume_reading(self, resume_parser: bool = True) -> None: 

199 was_paused = self._reading_paused 

200 super().resume_reading(resume_parser) 

201 if was_paused: 

202 self._reschedule_timeout() 

203 

204 def set_exception( 

205 self, 

206 exc: type[BaseException] | BaseException, 

207 exc_cause: BaseException = _EXC_SENTINEL, 

208 ) -> None: 

209 self._should_close = True 

210 self._drop_timeout() 

211 super().set_exception(exc, exc_cause) 

212 

213 def set_parser( 

214 self, 

215 parser: WebSocketReader, 

216 payload: WebSocketDataQueue, 

217 data_received_cb: Callable[[], None] | None = None, 

218 ) -> None: 

219 self._payload = payload 

220 self._payload_parser = parser 

221 self._data_received_cb = data_received_cb 

222 

223 self._drop_timeout() 

224 

225 if self._tail: 

226 data, self._tail = self._tail, b"" 

227 self.data_received(data) 

228 

229 def set_response_params( 

230 self, 

231 *, 

232 timer: BaseTimerContext | None = None, 

233 skip_payload: bool = False, 

234 read_until_eof: bool = False, 

235 auto_decompress: bool = True, 

236 read_timeout: float | None = None, 

237 read_bufsize: int = DEFAULT_CHUNK_SIZE, 

238 timeout_ceil_threshold: float = 5, 

239 max_line_size: int = 8190, 

240 max_field_size: int = 8190, 

241 max_headers: int = 128, 

242 ) -> None: 

243 self._skip_payload = skip_payload 

244 

245 self._read_timeout = read_timeout 

246 

247 self._timeout_ceil_threshold = timeout_ceil_threshold 

248 

249 self._parser = HttpResponseParser( 

250 self, 

251 self._loop, 

252 read_bufsize, 

253 timer=timer, 

254 payload_exception=ClientPayloadError, 

255 response_with_body=not skip_payload, 

256 read_until_eof=read_until_eof, 

257 auto_decompress=auto_decompress, 

258 max_line_size=max_line_size, 

259 max_field_size=max_field_size, 

260 max_headers=max_headers, 

261 ) 

262 

263 if self._tail: 

264 data, self._tail = self._tail, b"" 

265 self.data_received(data) 

266 

267 def _drop_timeout(self) -> None: 

268 if self._read_timeout_handle is not None: 

269 self._read_timeout_handle.cancel() 

270 self._read_timeout_handle = None 

271 

272 def _reschedule_timeout(self) -> None: 

273 timeout = self._read_timeout 

274 if self._read_timeout_handle is not None: 

275 self._read_timeout_handle.cancel() 

276 

277 if timeout: 

278 self._read_timeout_handle = self._loop.call_later( 

279 timeout, self._on_read_timeout 

280 ) 

281 else: 

282 self._read_timeout_handle = None 

283 

284 def start_timeout(self) -> None: 

285 self._reschedule_timeout() 

286 

287 @property 

288 def read_timeout(self) -> float | None: 

289 return self._read_timeout 

290 

291 @read_timeout.setter 

292 def read_timeout(self, read_timeout: float | None) -> None: 

293 self._read_timeout = read_timeout 

294 

295 def _on_read_timeout(self) -> None: 

296 exc = SocketTimeoutError("Timeout on reading data from socket") 

297 self.set_exception(exc) 

298 if self._payload is not None: 

299 set_exception(self._payload, exc) 

300 

301 def data_received(self, data: bytes) -> None: 

302 # If no data, then we are resuming decompression. We haven't received 

303 # data from the socket, so we can avoid the reschedule overhead. 

304 if data: 

305 self._reschedule_timeout() 

306 

307 # custom payload parser - currently always WebSocketReader 

308 if self._payload_parser is not None: 

309 if self._data_received_cb is not None: 

310 self._data_received_cb() 

311 eof, tail = self._payload_parser.feed_data(data) 

312 if eof: 

313 self._payload = None 

314 self._payload_parser = None 

315 

316 if tail: 

317 self.data_received(tail) 

318 return 

319 

320 if self._upgraded or self._parser is None: 

321 # i.e. websocket connection, websocket parser is not set yet 

322 self._tail += data 

323 return 

324 

325 # parse http messages 

326 try: 

327 messages, upgraded, tail = self._parser.feed_data(data) 

328 except BaseException as underlying_exc: 

329 if self.transport is not None: 

330 # connection.release() could be called BEFORE 

331 # data_received(), the transport is already 

332 # closed in this case 

333 self.transport.close() 

334 if not isinstance(underlying_exc, Exception): 

335 raise 

336 # should_close is True after the call 

337 if isinstance(underlying_exc, HttpProcessingError): 

338 exc = HttpProcessingError( 

339 code=underlying_exc.code, 

340 message=underlying_exc.message, 

341 headers=underlying_exc.headers, 

342 ) 

343 else: 

344 exc = HttpProcessingError() 

345 self.set_exception(exc, underlying_exc) 

346 return 

347 

348 self._upgraded = upgraded 

349 

350 payload: StreamReader | None = None 

351 for message, payload in messages: 

352 if message.should_close: 

353 self._should_close = True 

354 

355 self._payload = payload 

356 

357 if self._skip_payload or message.code in EMPTY_BODY_STATUS_CODES: 

358 self.feed_data((message, EMPTY_PAYLOAD)) 

359 else: 

360 self.feed_data((message, payload)) 

361 

362 if payload is not None: 

363 # new message(s) was processed 

364 # register timeout handler unsubscribing 

365 # either on end-of-stream or immediately for 

366 # EMPTY_PAYLOAD 

367 if payload is not EMPTY_PAYLOAD: 

368 payload.on_eof(self._drop_timeout) 

369 else: 

370 self._drop_timeout() 

371 

372 if upgraded and tail: 

373 self.data_received(tail)