Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/urllib3/util/wait.py: 22%

1from __future__ import annotations 

2 

3import select 

4import socket 

5from functools import partial 

6 

7__all__ = ["wait_for_read", "wait_for_write"] 

8 

9 

10# How should we wait on sockets? 

11# 

12# There are two types of APIs you can use for waiting on sockets: the fancy 

13# modern stateful APIs like epoll/kqueue, and the older stateless APIs like 

14# select/poll. The stateful APIs are more efficient when you have a lots of 

15# sockets to keep track of, because you can set them up once and then use them 

16# lots of times. But we only ever want to wait on a single socket at a time 

17# and don't want to keep track of state, so the stateless APIs are actually 

18# more efficient. So we want to use select() or poll(). 

19# 

20# Now, how do we choose between select() and poll()? On traditional Unixes, 

21# select() has a strange calling convention that makes it slow, or fail 

22# altogether, for high-numbered file descriptors. The point of poll() is to fix 

23# that, so on Unixes, we prefer poll(). 

24# 

25# On Windows, there is no poll() (or at least Python doesn't provide a wrapper 

26# for it), but that's OK, because on Windows, select() doesn't have this 

27# strange calling convention; plain select() works fine. 

28# 

29# So: on Windows we use select(), and everywhere else we use poll(). We also 

30# fall back to select() in case poll() is somehow broken or missing. 

31 

32 

33def select_wait_for_socket( 

34 sock: socket.socket, 

35 read: bool = False, 

36 write: bool = False, 

37 timeout: float | None = None, 

38) -> bool: 

39 if not read and not write: 

40 raise RuntimeError("must specify at least one of read=True, write=True") 

41 rcheck = [] 

42 wcheck = [] 

43 if read: 

44 rcheck.append(sock) 

45 if write: 

46 wcheck.append(sock) 

47 # When doing a non-blocking connect, most systems signal success by 

48 # marking the socket writable. Windows, though, signals success by marked 

49 # it as "exceptional". We paper over the difference by checking the write 

50 # sockets for both conditions. (The stdlib selectors module does the same 

51 # thing.) 

52 fn = partial(select.select, rcheck, wcheck, wcheck) 

53 rready, wready, xready = fn(timeout) 

54 return bool(rready or wready or xready) 

55 

56 

57def poll_wait_for_socket( 

58 sock: socket.socket, 

59 read: bool = False, 

60 write: bool = False, 

61 timeout: float | None = None, 

62) -> bool: 

63 if not read and not write: 

64 raise RuntimeError("must specify at least one of read=True, write=True") 

65 mask = 0 

66 if read: 

67 mask |= select.POLLIN 

68 if write: 

69 mask |= select.POLLOUT 

70 poll_obj = select.poll() 

71 poll_obj.register(sock, mask) 

72 

73 # For some reason, poll() takes timeout in milliseconds 

74 def do_poll(t: float | None) -> list[tuple[int, int]]: 

75 if t is not None: 

76 t *= 1000 

77 return poll_obj.poll(t) 

78 

79 return bool(do_poll(timeout)) 

80 

81 

82def _have_working_poll() -> bool: 

83 # Apparently some systems have a select.poll that fails as soon as you try 

84 # to use it, either due to strange configuration or broken monkeypatching 

85 # from libraries like eventlet/greenlet. 

86 try: 

87 poll_obj = select.poll() 

88 poll_obj.poll(0) 

89 except (AttributeError, OSError): 

90 return False 

91 else: 

92 return True 

93 

94 

95def wait_for_socket( 

96 sock: socket.socket, 

97 read: bool = False, 

98 write: bool = False, 

99 timeout: float | None = None, 

100) -> bool: 

101 # We delay choosing which implementation to use until the first time we're 

102 # called. We could do it at import time, but then we might make the wrong 

103 # decision if someone goes wild with monkeypatching select.poll after 

104 # we're imported. 

105 global wait_for_socket 

106 if _have_working_poll(): 

107 wait_for_socket = poll_wait_for_socket 

108 elif hasattr(select, "select"): 

109 wait_for_socket = select_wait_for_socket 

110 return wait_for_socket(sock, read, write, timeout) 

111 

112 

113def wait_for_read(sock: socket.socket, timeout: float | None = None) -> bool: 

114 """Waits for reading to be available on a given socket. 

115 Returns True if the socket is readable, or False if the timeout expired. 

116 """ 

117 return wait_for_socket(sock, read=True, timeout=timeout) 

118 

119 

120def wait_for_write(sock: socket.socket, timeout: float | None = None) -> bool: 

121 """Waits for writing to be available on a given socket. 

122 Returns True if the socket is readable, or False if the timeout expired. 

123 """ 

124 return wait_for_socket(sock, write=True, timeout=timeout)