codekingpro/portable-devtools

1from __future__ import annotations2 3import select4import socket5from functools import partial6 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 fancy13# modern stateful APIs like epoll/kqueue, and the older stateless APIs like14# select/poll. The stateful APIs are more efficient when you have a lots of15# sockets to keep track of, because you can set them up once and then use them16# lots of times. But we only ever want to wait on a single socket at a time17# and don't want to keep track of state, so the stateless APIs are actually18# 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 fail22# altogether, for high-numbered file descriptors. The point of poll() is to fix23# that, so on Unixes, we prefer poll().24#25# On Windows, there is no poll() (or at least Python doesn't provide a wrapper26# for it), but that's OK, because on Windows, select() doesn't have this27# strange calling convention; plain select() works fine.28#29# So: on Windows we use select(), and everywhere else we use poll(). We also30# 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 by48    # marking the socket writable. Windows, though, signals success by marked49    # it as "exceptional". We paper over the difference by checking the write50    # sockets for both conditions. (The stdlib selectors module does the same51    # 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 = 066    if read:67        mask |= select.POLLIN68    if write:69        mask |= select.POLLOUT70    poll_obj = select.poll()71    poll_obj.register(sock, mask)72 73    # For some reason, poll() takes timeout in milliseconds74    def do_poll(t: float | None) -> list[tuple[int, int]]:75        if t is not None:76            t *= 100077        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 try84    # to use it, either due to strange configuration or broken monkeypatching85    # from libraries like eventlet/greenlet.86    try:87        poll_obj = select.poll()88        poll_obj.poll(0)89    except (AttributeError, OSError):90        return False91    else:92        return True93 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're102    # called. We could do it at import time, but then we might make the wrong103    # decision if someone goes wild with monkeypatching select.poll after104    # we're imported.105    global wait_for_socket106    if _have_working_poll():107        wait_for_socket = poll_wait_for_socket108    elif hasattr(select, "select"):109        wait_for_socket = select_wait_for_socket110    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)125