Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/urllib3/util/wait.py: 22%
49 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-08 06:05 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-08 06:05 +0000
1from __future__ import annotations
3import select
4import socket
5from functools import partial
7__all__ = ["wait_for_read", "wait_for_write"]
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.
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)
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)
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)
79 return bool(do_poll(timeout))
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
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)
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)
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)