Merge pull request #610 from njsmith/fix-400

wait_socket_*, notify_socket_close now accept raw socket descriptors

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

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.

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"
     ]

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)

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.