Replace asyncio.wait_for with asyncio.timeout.

Some instances were missed in a previous commit.

Also update documentation.

Author: aaugustin

docs/faq/common.rst

@@ -105,9 +105,14 @@ you can adjust it with the ``ping_timeout`` argument.
 How do I set a timeout on :meth:`~legacy.protocol.WebSocketCommonProtocol.recv`?
 --------------------------------------------------------------------------------
 
-Use :func:`~asyncio.wait_for`::
+On Python ≥ 3.11, use :func:`asyncio.timeout`::
 
-    await asyncio.wait_for(websocket.recv(), timeout=10)
+    async with asyncio.timeout(timeout=10):
+        message = await websocket.recv()
+
+On older versions of Python, use :func:`asyncio.wait_for`::
+
+    message = await asyncio.wait_for(websocket.recv(), timeout=10)
 
 This technique works for most APIs. When it doesn't, for example with
 asynchronous context managers, websockets provides an ``open_timeout`` argument.

docs/topics/design.rst

@@ -437,7 +437,7 @@ propagate cancellation to them.
 prevent cancellation.
 
 :meth:`~legacy.protocol.WebSocketCommonProtocol.close` waits for the data transfer
-task to terminate with :func:`~asyncio.wait_for`. If it's canceled or if the
+task to terminate with :func:`~asyncio.timeout`. If it's canceled or if the
 timeout elapses, :attr:`~legacy.protocol.WebSocketCommonProtocol.transfer_data_task`
 is canceled, which is correct at this point.
 :meth:`~legacy.protocol.WebSocketCommonProtocol.close` then waits for

src/websockets/legacy/client.py

@@ -44,6 +44,7 @@
 from ..http import USER_AGENT
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
 from ..uri import WebSocketURI, parse_uri
+from .compatibility import asyncio_timeout
 from .handshake import build_request, check_response
 from .http import read_response
 from .protocol import WebSocketCommonProtocol
@@ -650,7 +651,8 @@ def __await__(self) -> Generator[Any, None, WebSocketClientProtocol]:
         return self.__await_impl_timeout__().__await__()
 
     async def __await_impl_timeout__(self) -> WebSocketClientProtocol:
-        return await asyncio.wait_for(self.__await_impl__(), self.open_timeout)
+        async with asyncio_timeout(self.open_timeout):
+            return await self.__await_impl__()
 
     async def __await_impl__(self) -> WebSocketClientProtocol:
         for redirects in range(self.MAX_REDIRECTS_ALLOWED):

src/websockets/legacy/protocol.py

@@ -107,7 +107,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
       context manager;
     * on the server side, when the connection handler terminates.
 
-    To apply a timeout to any other API, wrap it in :func:`~asyncio.wait_for`.
+    To apply a timeout to any other API, wrap it in :func:`~asyncio.timeout` or
+    :func:`~asyncio.wait_for`.
 
     The ``max_size`` parameter enforces the maximum size for incoming messages
     in bytes. The default value is 1 MiB. If a larger message is received,
@@ -513,7 +514,7 @@ async def recv(self) -> Data:
         message. The next invocation of :meth:`recv` will return it.
 
         This makes it possible to enforce a timeout by wrapping :meth:`recv` in
-        :func:`~asyncio.wait_for`.
+        :func:`~asyncio.timeout` or :func:`~asyncio.wait_for`.
 
         Returns:
             Data: A string (:class:`str`) for a Text_ frame. A bytestring

src/websockets/legacy/server.py

@@ -46,7 +46,7 @@
 from ..http import USER_AGENT
 from ..protocol import State
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
-from .compatibility import loop_if_py_lt_38
+from .compatibility import asyncio_timeout, loop_if_py_lt_38
 from .handshake import build_response, check_request
 from .http import read_request
 from .protocol import WebSocketCommonProtocol
@@ -163,15 +163,13 @@ async def handler(self) -> None:
         """
         try:
             try:
-                await asyncio.wait_for(
-                    self.handshake(
+                async with asyncio_timeout(self.open_timeout):
+                    await self.handshake(
                         origins=self.origins,
                         available_extensions=self.available_extensions,
                         available_subprotocols=self.available_subprotocols,
                         extra_headers=self.extra_headers,
-                    ),
-                    self.open_timeout,
-                )
+                    )
             # Remove this branch when dropping support for Python < 3.8
             # because CancelledError no longer inherits Exception.
             except asyncio.CancelledError:  # pragma: no cover

tests/legacy/test_client_server.py

@@ -29,6 +29,7 @@
 )
 from websockets.http import USER_AGENT
 from websockets.legacy.client import *
+from websockets.legacy.compatibility import asyncio_timeout
 from websockets.legacy.handshake import build_response
 from websockets.legacy.http import read_response
 from websockets.legacy.server import *
@@ -1129,7 +1130,8 @@ def test_client_connect_canceled_during_handshake(self):
 
         async def cancelled_client():
             start_client = connect(get_server_uri(self.server), sock=sock)
-            await asyncio.wait_for(start_client, 5 * MS)
+            async with asyncio_timeout(5 * MS):
+                await start_client
 
         with self.assertRaises(asyncio.TimeoutError):
             self.loop.run_until_complete(cancelled_client())