Move asyncio compatibility to a new package.

Author: aaugustin

src/websockets/asyncio/__init__.py

@@ -20,6 +20,7 @@
     cast,
 )
 
+from ..asyncio.compatibility import asyncio_timeout
 from ..datastructures import Headers, HeadersLike
 from ..exceptions import (
     InvalidHandshake,
@@ -44,7 +45,6 @@
 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

src/websockets/legacy/async_timeout.py src/websockets/asyncio/async_timeout.py

@@ -28,6 +28,7 @@
     cast,
 )
 
+from ..asyncio.compatibility import asyncio_timeout
 from ..datastructures import Headers
 from ..exceptions import (
     ConnectionClosed,
@@ -53,7 +54,6 @@
 )
 from ..protocol import State
 from ..typing import Data, LoggerLike, Subprotocol
-from .compatibility import asyncio_timeout
 from .framing import Frame
 
 

src/websockets/legacy/compatibility.py src/websockets/asyncio/compatibility.py

@@ -25,6 +25,7 @@
     cast,
 )
 
+from ..asyncio.compatibility import asyncio_timeout
 from ..datastructures import Headers, HeadersLike, MultipleValuesError
 from ..exceptions import (
     AbortHandshake,
@@ -46,7 +47,6 @@
 from ..http import USER_AGENT
 from ..protocol import State
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
-from .compatibility import asyncio_timeout
 from .handshake import build_response, check_request
 from .http import read_request
 from .protocol import WebSocketCommonProtocol

src/websockets/legacy/client.py

@@ -72,8 +72,10 @@ def get(self, timeout: Optional[float] = None) -> Data:
 
         Raises:
             EOFError: If the stream of frames has ended.
-            RuntimeError: If two threads run :meth:`get` or :meth:``get_iter`
+            RuntimeError: If two threads run :meth:`get` or :meth:`get_iter`
                 concurrently.
+            TimeoutError: If a timeout is provided and elapses before a
+                complete message is received.
 
         """
         with self.mutex:
@@ -131,7 +133,7 @@ def get_iter(self) -> Iterator[Data]:
 
         Raises:
             EOFError: If the stream of frames has ended.
-            RuntimeError: If two threads run :meth:`get` or :meth:``get_iter`
+            RuntimeError: If two threads run :meth:`get` or :meth:`get_iter`
                 concurrently.
 
         """
@@ -159,11 +161,9 @@ def get_iter(self) -> Iterator[Data]:
             self.get_in_progress = True
 
         # Locking with get_in_progress ensures only one thread can get here.
-        yield from chunks
-        while True:
-            chunk = self.chunks_queue.get()
-            if chunk is None:
-                break
+        for chunk in chunks:
+            yield chunk
+        while (chunk := self.chunks_queue.get()) is not None:
             yield chunk
 
         with self.mutex:
@@ -242,6 +242,7 @@ def put(self, frame: Frame) -> None:
             self.put_in_progress = True
 
         # Release the lock to allow get() to run and eventually set the event.
+        # Locking with get_in_progress ensures only one coroutine can get here.
         self.message_fetched.wait()
 
         with self.mutex:

src/websockets/legacy/protocol.py

@@ -14,6 +14,7 @@
 import urllib.request
 import warnings
 
+from websockets.asyncio.compatibility import asyncio_timeout
 from websockets.datastructures import Headers
 from websockets.exceptions import (
     ConnectionClosed,
@@ -29,7 +30,6 @@
 )
 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 *