diff --git a/docs/source/reference-hazmat.rst b/docs/source/reference-hazmat.rst index efee0ba528..054a22f7dc 100644 --- a/docs/source/reference-hazmat.rst +++ b/docs/source/reference-hazmat.rst @@ -129,11 +129,11 @@ All environments provide the following functions: Block until the given :func:`socket.socket` object is readable. - The given object *must* be exactly of type :func:`socket.socket`, - nothing else. + On Unix systems, sockets are fds, and this is identical to + :func:`wait_readable`. On Windows, ``SOCKET`` handles and fds are + different, and this works on ``SOCKET`` handles or Python socket + objects. - :raises TypeError: - if the given object is not of type :func:`socket.socket`. :raises trio.ResourceBusyError: if another task is already waiting for the given socket to become readable. @@ -143,11 +143,11 @@ All environments provide the following functions: Block until the given :func:`socket.socket` object is writable. - The given object *must* be exactly of type :func:`socket.socket`, - nothing else. + On Unix systems, sockets are fds, and this is identical to + :func:`wait_writable`. On Windows, ``SOCKET`` handles and fds are + different, and this works on ``SOCKET`` handles or Python socket + objects. - :raises TypeError: - if the given object is not of type :func:`socket.socket`. :raises trio.ResourceBusyError: if another task is already waiting for the given socket to become writable. @@ -168,11 +168,10 @@ All environments provide the following functions: socket, you should first call this function, and then close the socket. - The given object *must* be exactly of type :func:`socket.socket`, - nothing else. - - :raises TypeError: - if the given object is not of type :func:`socket.socket`. + On Unix systems, sockets are fds, and this is identical to + :func:`notify_fd_close`. On Windows, ``SOCKET`` handles and fds are + different, and this works on ``SOCKET`` handles or Python socket + objects. Unix-specific API diff --git a/newsfragments/400.feature.rst b/newsfragments/400.feature.rst new file mode 100644 index 0000000000..e76e672e02 --- /dev/null +++ b/newsfragments/400.feature.rst @@ -0,0 +1,4 @@ +The low level :func:`trio.hazmat.wait_socket_readable`, +:func:`~trio.hazmat.wait_socket_writable`, and +:func:`~trio.hazmat.notify_socket_close` now work on bare socket +descriptors, instead of requiring a :func:`socket.socket` object. diff --git a/trio/_core/__init__.py b/trio/_core/__init__.py index 976975ae48..9dc27b971e 100644 --- a/trio/_core/__init__.py +++ b/trio/_core/__init__.py @@ -44,23 +44,9 @@ def _public(fn): __all__ += _local.__all__ if hasattr(_run, "wait_readable"): - import socket as _stdlib_socket - - async def wait_socket_readable(sock): - if type(sock) != _stdlib_socket.socket: - raise TypeError("need a socket") - await wait_readable(sock) - - async def wait_socket_writable(sock): - if type(sock) != _stdlib_socket.socket: - raise TypeError("need a socket") - await wait_writable(sock) - - def notify_socket_close(sock): - if type(sock) != _stdlib_socket.socket: - raise TypeError("need a socket") - notify_fd_close(sock) - + wait_socket_readable = wait_readable + wait_socket_writable = wait_writable + notify_socket_close = notify_fd_close __all__ += [ "wait_socket_readable", "wait_socket_writable", "notify_socket_close" ] diff --git a/trio/_core/_io_windows.py b/trio/_core/_io_windows.py index 885ddccf6f..d044b1f0a5 100644 --- a/trio/_core/_io_windows.py +++ b/trio/_core/_io_windows.py @@ -336,18 +336,8 @@ def monitor_completion_key(self): del self._completion_key_queues[key] async def _wait_socket(self, which, sock): - # Using socket objects rather than raw handles gives better behavior - # if someone closes the socket while another task is waiting on it. If - # we just kept the handle, it might be reassigned, and we'd be waiting - # on who-knows-what. The socket object won't be reassigned, and it - # switches its fileno() to -1, so we can detect the offending socket - # and wake the appropriate task. This is a pretty minor benefit (I - # think it can only make a difference if someone is closing random - # sockets in another thread? And on unix we don't handle this case at - # all), but hey, why not. - if type(sock) is not stdlib_socket.socket: - await _core.checkpoint() - raise TypeError("need a stdlib socket") + if not isinstance(sock, int): + sock = sock.fileno() if sock in self._socket_waiters[which]: await _core.checkpoint() raise _core.ResourceBusyError( @@ -372,9 +362,8 @@ async def wait_socket_writable(self, sock): @_public def notify_socket_close(self, sock): - if type(sock) is not stdlib_socket.socket: - raise TypeError("need a stdlib socket") - + if not isinstance(sock, int): + sock = sock.fileno() for mode in ["read", "write"]: if sock in self._socket_waiters[mode]: task = self._socket_waiters[mode].pop(sock) diff --git a/trio/_core/tests/test_io.py b/trio/_core/tests/test_io.py index dad67ab8fc..76abebc773 100644 --- a/trio/_core/tests/test_io.py +++ b/trio/_core/tests/test_io.py @@ -36,35 +36,37 @@ def socketpair(): sock.close() +def using_fileno(fn): + def fileno_wrapper(fileobj): + return fn(fileobj.fileno()) + + name = "<{} on fileno>".format(fn.__name__) + fileno_wrapper.__name__ = fileno_wrapper.__qualname__ = name + return fileno_wrapper + + wait_readable_options = [_core.wait_socket_readable] wait_writable_options = [_core.wait_socket_writable] notify_close_options = [_core.notify_socket_close] -if hasattr(_core, "wait_readable"): - wait_readable_options.append(_core.wait_readable) - async def wait_readable_fd(fileobj): - return await _core.wait_readable(fileobj.fileno()) - - wait_readable_options.append(wait_readable_fd) +# We aren't testing the _fd_ versions, because they're the same as the socket +# ones. But if they ever stop being the same we should notice and add tests! +if hasattr(_core, "wait_readable"): + assert _core.wait_socket_readable is _core.wait_readable if hasattr(_core, "wait_writable"): - wait_writable_options.append(_core.wait_writable) - - async def wait_writable_fd(fileobj): - return await _core.wait_writable(fileobj.fileno()) - - wait_writable_options.append(wait_writable_fd) + assert _core.wait_socket_writable is _core.wait_writable if hasattr(_core, "notify_fd_close"): - notify_close_options.append(_core.notify_fd_close) - - def notify_fd_close_rawfd(fileobj): - _core.notify_fd_close(fileobj.fileno()) + assert _core.notify_socket_close is _core.notify_fd_close - notify_close_options.append(notify_fd_close_rawfd) +for options_list in [ + wait_readable_options, wait_writable_options, notify_close_options +]: + options_list += [using_fileno(f) for f in options_list] # Decorators that feed in different settings for wait_readable / wait_writable # / notify_close. # Note that if you use all three decorators on the same test, it will run all -# N**3 *combinations* +# N**4 *combinations* read_socket_test = pytest.mark.parametrize( "wait_readable", wait_readable_options, ids=lambda fn: fn.__name__ ) @@ -76,26 +78,6 @@ def notify_fd_close_rawfd(fileobj): ) -async def test_wait_socket_type_checking(socketpair): - a, b = socketpair - - # wait_socket_* accept actual socket objects, only - for sock_fn in [ - _core.wait_socket_readable, - _core.wait_socket_writable, - _core.notify_socket_close, - ]: - with pytest.raises(TypeError): - await sock_fn(a.fileno()) - - class AllegedSocket(stdlib_socket.socket): - pass - - with AllegedSocket() as alleged_socket: - with pytest.raises(TypeError): - await sock_fn(alleged_socket) - - # XX These tests are all a bit dicey because they can't distinguish between # wait_on_{read,writ}able blocking the way it should, versus blocking # momentarily and then immediately resuming.