diff --git a/src/websockets/asyncio/__init__.py b/src/websockets/asyncio/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/websockets/legacy/async_timeout.py b/src/websockets/asyncio/async_timeout.py similarity index 100% rename from src/websockets/legacy/async_timeout.py rename to src/websockets/asyncio/async_timeout.py diff --git a/src/websockets/legacy/compatibility.py b/src/websockets/asyncio/compatibility.py similarity index 100% rename from src/websockets/legacy/compatibility.py rename to src/websockets/asyncio/compatibility.py diff --git a/src/websockets/legacy/client.py b/src/websockets/legacy/client.py index 48622523..586ba528 100644 --- a/src/websockets/legacy/client.py +++ b/src/websockets/legacy/client.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 diff --git a/src/websockets/legacy/protocol.py b/src/websockets/legacy/protocol.py index 0422c10d..9d2fe6e7 100644 --- a/src/websockets/legacy/protocol.py +++ b/src/websockets/legacy/protocol.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 diff --git a/src/websockets/legacy/server.py b/src/websockets/legacy/server.py index a17c5232..23322050 100644 --- a/src/websockets/legacy/server.py +++ b/src/websockets/legacy/server.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 diff --git a/src/websockets/sync/messages.py b/src/websockets/sync/messages.py index 67a22313..114a9f52 100644 --- a/src/websockets/sync/messages.py +++ b/src/websockets/sync/messages.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: diff --git a/tests/legacy/test_client_server.py b/tests/legacy/test_client_server.py index 133af053..0c96e821 100644 --- a/tests/legacy/test_client_server.py +++ b/tests/legacy/test_client_server.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 *