Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/base_protocol.py: 44%

1import asyncio 

2from typing import TYPE_CHECKING, Any, cast 

3 

4from .client_exceptions import ClientConnectionResetError 

5from .helpers import set_exception 

6from .tcp_helpers import tcp_nodelay 

7 

8if TYPE_CHECKING: 

9 from .http_parser import HttpParser 

10 

11# Raised by transport.pause_reading()/resume_reading() when the transport 

12# does not support flow control; safe to ignore. 

13# NOTE: Catch these with a plain try/except/pass, never contextlib.suppress(): 

14# pause/resume run on the hot read path and suppress() is ~6x slower than 

15# try/except here (it builds a context manager and unpacks this tuple per call). 

16PAUSE_RESUME_READING_ERRORS = (AttributeError, NotImplementedError, RuntimeError) 

17 

18 

19class BaseProtocol(asyncio.Protocol): 

20 __slots__ = ( 

21 "_loop", 

22 "_paused", 

23 "_parser", 

24 "_drain_waiter", 

25 "_connection_lost", 

26 "_reading_paused", 

27 "_buffer_paused", 

28 "_upgraded", 

29 "transport", 

30 ) 

31 

32 def __init__( 

33 self, loop: asyncio.AbstractEventLoop, parser: "HttpParser[Any] | None" = None 

34 ) -> None: 

35 self._loop: asyncio.AbstractEventLoop = loop 

36 self._paused = False 

37 self._drain_waiter: asyncio.Future[None] | None = None 

38 self._reading_paused = False 

39 self._buffer_paused = False 

40 self._parser = parser 

41 self._upgraded = False 

42 

43 self.transport: asyncio.Transport | None = None 

44 

45 @property 

46 def connected(self) -> bool: 

47 """Return True if the connection is open.""" 

48 return self.transport is not None 

49 

50 @property 

51 def writing_paused(self) -> bool: 

52 return self._paused 

53 

54 def pause_writing(self) -> None: 

55 assert not self._paused 

56 self._paused = True 

57 

58 def resume_writing(self) -> None: 

59 assert self._paused 

60 self._paused = False 

61 

62 waiter = self._drain_waiter 

63 if waiter is not None: 

64 self._drain_waiter = None 

65 if not waiter.done(): 

66 waiter.set_result(None) 

67 

68 def pause_reading(self) -> None: 

69 self._reading_paused = True 

70 # Parser shouldn't be paused on websockets. 

71 if not self._upgraded: 

72 assert self._parser is not None 

73 self._parser.pause_reading() 

74 self._pause_transport_reading() 

75 

76 def _pause_reading_for_buffer(self) -> None: 

77 """Hold the transport for a buffer this protocol cannot drain yet. 

78 

79 One flag covers every buffer a protocol owns; the server already 

80 multiplexes its message queue and its tail through it. Re-check the 

81 conditions on resume rather than adding another flag. 

82 """ 

83 self._buffer_paused = True 

84 self._pause_transport_reading() 

85 

86 def _resume_reading_for_buffer(self) -> None: 

87 """Release that hold, unless flow control is still holding it too.""" 

88 self._buffer_paused = False 

89 if not self._reading_paused: 

90 self._resume_transport_reading() 

91 

92 def _pause_transport_reading(self) -> None: 

93 if self.transport is not None: 

94 try: 

95 self.transport.pause_reading() 

96 except PAUSE_RESUME_READING_ERRORS: 

97 # Transport lacks flow control; nothing to pause. Intentionally 

98 # ignored (see PAUSE_RESUME_READING_ERRORS; do not use suppress). 

99 pass 

100 

101 def _resume_transport_reading(self) -> None: 

102 if self.transport is not None: 

103 try: 

104 self.transport.resume_reading() 

105 except PAUSE_RESUME_READING_ERRORS: 

106 # Transport lacks flow control; nothing to resume. Intentionally 

107 # ignored (see PAUSE_RESUME_READING_ERRORS; do not use suppress). 

108 pass 

109 

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

111 self._reading_paused = False 

112 

113 # This will resume parsing any unprocessed data from the last pause. 

114 if not self._upgraded and resume_parser: 

115 self.data_received(b"") 

116 

117 # Reading may have been paused again in the above call if there was a lot of 

118 # compressed data still pending. 

119 if not self._reading_paused and not self._buffer_paused: 

120 self._resume_transport_reading() 

121 

122 def connection_made(self, transport: asyncio.BaseTransport) -> None: 

123 tr = cast(asyncio.Transport, transport) 

124 tcp_nodelay(tr, True) 

125 self.transport = tr 

126 

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

128 # Wake up the writer if currently paused. 

129 self.transport = None 

130 if not self._paused: 

131 return 

132 waiter = self._drain_waiter 

133 if waiter is None: 

134 return 

135 self._drain_waiter = None 

136 if waiter.done(): 

137 return 

138 if exc is None: 

139 waiter.set_result(None) 

140 else: 

141 set_exception( 

142 waiter, 

143 ConnectionError("Connection lost"), 

144 exc, 

145 ) 

146 

147 async def _drain_helper(self) -> None: 

148 if self.transport is None: 

149 raise ClientConnectionResetError("Connection lost") 

150 if not self._paused: 

151 return 

152 waiter = self._drain_waiter 

153 if waiter is None: 

154 waiter = self._loop.create_future() 

155 self._drain_waiter = waiter 

156 await waiter