Skip to content

Add trio implementation - step 3 #1673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aaugustin wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add trio implementation - step 3 #1673

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
175d297
Align the asyncio and sync client and server modules.
aaugustin Nov 9, 2025
1d34ddb
Clean up the asyncio and sync client and server tests.
aaugustin Nov 9, 2025
aa15fc4
Add trio client and server.
aaugustin Nov 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
256 changes: 128 additions & 128 deletions docs/reference/features.rst
View file
Open in desktop

Large diffs are not rendered by default.

156 changes: 87 additions & 69 deletions src/websockets/asyncio/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,8 +214,8 @@ class connect:
compression: The "permessage-deflate" extension is enabled by default.
Set ``compression`` to :obj:`None` to disable it. See the
:doc:`compression guide <../../topics/compression>` for details.
additional_headers (HeadersLike | None): Arbitrary HTTP headers to add
to the handshake request.
additional_headers: Arbitrary HTTP headers to add to the handshake
request.
user_agent_header: Value of the ``User-Agent`` request header.
It defaults to ``"Python/x.y.z websockets/X.Y"``.
Setting it to :obj:`None` removes the header.
Expand Down Expand Up @@ -261,23 +261,25 @@ class connect:

For example:

* You can set ``sock`` to provide a preexisting TCP socket. You may call
:func:`socket.create_connection` (not to be confused with the event loop's
:meth:`~asyncio.loop.create_connection` method) to create a suitable
client socket and customize it.

* You can set ``ssl`` to a :class:`~ssl.SSLContext` to enforce TLS settings.
When connecting to a ``wss://`` URI, if ``ssl`` isn't provided, a TLS
context is created with :func:`~ssl.create_default_context`.

* You can set ``server_hostname`` to override the host name from ``uri`` in
the TLS handshake.

* You can configure ``ssl_handshake_timeout``, and ``ssl_shutdown_timeout``.

* You can set ``host`` and ``port`` to connect to a different host and port
from those found in ``uri``. This only changes the destination of the TCP
connection. The host name from ``uri`` is still used in the TLS handshake
for secure connections and in the ``Host`` header.

* You can set ``sock`` to provide a preexisting TCP socket. You may call
:func:`socket.create_connection` (not to be confused with the event loop's
:meth:`~asyncio.loop.create_connection` method) to create a suitable
client socket and customize it.

When using a proxy:

* Prefix keyword arguments with ``proxy_`` for configuring TLS between the
Expand Down Expand Up @@ -329,6 +331,9 @@ def __init__(
**kwargs: Any,
) -> None:
self.uri = uri
self.ws_uri = parse_uri(uri)
if not self.ws_uri.secure and kwargs.get("ssl") is not None:
raise ValueError("ssl argument is incompatible with a ws:// URI")

if subprotocols is not None:
validate_subprotocols(subprotocols)
Expand All @@ -344,7 +349,16 @@ def __init__(
if create_connection is None:
create_connection = ClientConnection

def protocol_factory(uri: WebSocketURI) -> ClientConnection:
self.additional_headers = additional_headers
self.user_agent_header = user_agent_header
self.proxy = proxy
self.process_exception = process_exception
self.open_timeout = open_timeout
self.reconnect_delays = reconnect_delays
self.logger = logger
self.create_connection_kwargs = kwargs

def factory(uri: WebSocketURI) -> ClientConnection:
# This is a protocol in the Sans-I/O implementation of websockets.
protocol = ClientProtocol(
uri,
Expand All @@ -355,6 +369,7 @@ def protocol_factory(uri: WebSocketURI) -> ClientConnection:
logger=logger,
)
# This is a connection in websockets and a protocol in asyncio.
# create_connection defaults to ClientConnection.
connection = create_connection(
protocol,
ping_interval=ping_interval,
Expand All @@ -365,52 +380,43 @@ def protocol_factory(uri: WebSocketURI) -> ClientConnection:
)
return connection

self.proxy = proxy
self.protocol_factory = protocol_factory
self.additional_headers = additional_headers
self.user_agent_header = user_agent_header
self.process_exception = process_exception
self.open_timeout = open_timeout
self.reconnect_delays = reconnect_delays
self.logger = logger
self.connection_kwargs = kwargs
self.factory = factory

async def create_connection(self) -> ClientConnection:
"""Create TCP or Unix connection."""
async def open_tcp_connection(self) -> ClientConnection:
"""Open a TCP or Unix connection to the server, possibly through a proxy."""
loop = asyncio.get_running_loop()
kwargs = self.connection_kwargs.copy()

ws_uri = parse_uri(self.uri)
kwargs = self.create_connection_kwargs.copy()

proxy = self.proxy
if kwargs.get("unix", False):
proxy = None
if kwargs.get("sock") is not None:
proxy = None
if proxy is True:
proxy = get_proxy(ws_uri)
proxy = get_proxy(self.ws_uri)

def factory() -> ClientConnection:
return self.protocol_factory(ws_uri)
return self.factory(self.ws_uri)

if ws_uri.secure:
if self.ws_uri.secure:
kwargs.setdefault("ssl", True)
kwargs.setdefault("server_hostname", ws_uri.host)
if kwargs.get("ssl") is None:
raise ValueError("ssl=None is incompatible with a wss:// URI")
else:
if kwargs.get("ssl") is not None:
raise ValueError("ssl argument is incompatible with a ws:// URI")
kwargs.setdefault("server_hostname", self.ws_uri.host)

if kwargs.pop("unix", False):
_, connection = await loop.create_unix_connection(factory, **kwargs)
elif proxy is not None:
return connection

if proxy is not None:
proxy_parsed = parse_proxy(proxy)

if proxy_parsed.scheme[:5] == "socks":
# Connect to the server through the proxy.
# Connect to the server through the SOCKS proxy.
sock = await connect_socks_proxy(
proxy_parsed,
ws_uri,
self.ws_uri,
# python_socks uses local_addr like asyncio.
local_addr=kwargs.pop("local_addr", None),
)
# Initialize WebSocket connection via the proxy.
Expand All @@ -419,6 +425,8 @@ def factory() -> ClientConnection:
sock=sock,
**kwargs,
)
return connection

elif proxy_parsed.scheme[:4] == "http":
# Split keyword arguments between the proxy and the server.
all_kwargs, proxy_kwargs, kwargs = kwargs, {}, {}
Expand All @@ -441,10 +449,10 @@ def factory() -> ClientConnection:
raise ValueError(
"proxy_ssl argument is incompatible with an http:// proxy"
)
# Connect to the server through the proxy.
# Connect to the server through the HTTP proxy.
transport = await connect_http_proxy(
proxy_parsed,
ws_uri,
self.ws_uri,
user_agent_header=self.user_agent_header,
**proxy_kwargs,
)
Expand All @@ -461,16 +469,19 @@ def factory() -> ClientConnection:
assert new_transport is not None # help mypy
transport = new_transport
connection.connection_made(transport)
return connection

else:
raise AssertionError("unsupported proxy")
else:
raise AssertionError("parse_proxy returned unsupported proxy")

else: # proxy is None
# Connect to the server directly.
if kwargs.get("sock") is None:
kwargs.setdefault("host", ws_uri.host)
kwargs.setdefault("port", ws_uri.port)
kwargs.setdefault("host", self.ws_uri.host)
kwargs.setdefault("port", self.ws_uri.port)
# Initialize WebSocket connection.
_, connection = await loop.create_connection(factory, **kwargs)
return connection
return connection

def process_redirect(self, exc: Exception) -> Exception | str:
"""
Expand All @@ -494,12 +505,12 @@ def process_redirect(self, exc: Exception) -> Exception | str:
):
return exc

old_ws_uri = parse_uri(self.uri)
old_ws_uri = self.ws_uri
new_uri = urllib.parse.urljoin(self.uri, exc.response.headers["Location"])
new_ws_uri = parse_uri(new_uri)

# If connect() received a socket, it is closed and cannot be reused.
if self.connection_kwargs.get("sock") is not None:
if self.create_connection_kwargs.get("sock") is not None:
return ValueError(
f"cannot follow redirect to {new_uri} with a preexisting socket"
)
Expand All @@ -515,16 +526,16 @@ def process_redirect(self, exc: Exception) -> Exception | str:
or old_ws_uri.port != new_ws_uri.port
):
# Cross-origin redirects on Unix sockets don't quite make sense.
if self.connection_kwargs.get("unix", False):
if self.create_connection_kwargs.get("unix", False):
return ValueError(
f"cannot follow cross-origin redirect to {new_uri} "
f"with a Unix socket"
)

# Cross-origin redirects when host and port are overridden are ill-defined.
if (
self.connection_kwargs.get("host") is not None
or self.connection_kwargs.get("port") is not None
self.create_connection_kwargs.get("host") is not None
or self.create_connection_kwargs.get("port") is not None
):
return ValueError(
f"cannot follow cross-origin redirect to {new_uri} "
Expand Down Expand Up @@ -554,14 +565,14 @@ async def __await_impl__(self) -> ClientConnection:
try:
async with asyncio.timeout(self.open_timeout):
for _ in range(MAX_REDIRECTS):
self.connection = await self.create_connection()
connection = await self.open_tcp_connection()
try:
await self.connection.handshake(
await connection.handshake(
self.additional_headers,
self.user_agent_header,
)
except asyncio.CancelledError:
self.connection.transport.abort()
connection.transport.abort()
raise
except Exception as exc:
# Always close the connection even though keep-alive is
Expand All @@ -570,22 +581,24 @@ async def __await_impl__(self) -> ClientConnection:
# protocol. In the current design of connect(), there is
# no easy way to reuse the network connection that works
# in every case nor to reinitialize the protocol.
self.connection.transport.abort()

uri_or_exc = self.process_redirect(exc)
# Response is a valid redirect; follow it.
if isinstance(uri_or_exc, str):
self.uri = uri_or_exc
continue
# Response isn't a valid redirect; raise the exception.
if uri_or_exc is exc:
raise
connection.transport.abort()

exc_or_uri = self.process_redirect(exc)
if isinstance(exc_or_uri, Exception):
# Response isn't a valid redirect; raise the exception.
if exc_or_uri is exc:
raise
else:
raise exc_or_uri from exc
else:
raise uri_or_exc from exc
# Response is a valid redirect; follow it.
self.uri = exc_or_uri
self.ws_uri = parse_uri(exc_or_uri)
continue

else:
self.connection.start_keepalive()
return self.connection
connection.start_keepalive()
return connection
else:
raise SecurityError(f"more than {MAX_REDIRECTS} redirects")

Expand All @@ -596,24 +609,30 @@ async def __await_impl__(self) -> ClientConnection:
# async with connect(...) as ...: ...

async def __aenter__(self) -> ClientConnection:
return await self
if hasattr(self, "connection"):
raise RuntimeError("connect() isn't reentrant")
self.connection = await self
return self.connection

async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
await self.connection.close()
try:
await self.connection.close()
finally:
del self.connection

# async for ... in connect(...):

async def __aiter__(self) -> AsyncIterator[ClientConnection]:
delays: Generator[float] | None = None
while True:
try:
async with self as protocol:
yield protocol
async with self as connection:
yield connection
except Exception as exc:
# Determine whether the exception is retryable or fatal.
# The API of process_exception is "return an exception or None";
Expand Down Expand Up @@ -642,7 +661,6 @@ async def __aiter__(self) -> AsyncIterator[ClientConnection]:
traceback.format_exception_only(exc)[0].strip(),
)
await asyncio.sleep(delay)
continue

else:
# The connection succeeded. Reset backoff.
Expand Down Expand Up @@ -786,13 +804,13 @@ def eof_received(self) -> None:

def connection_lost(self, exc: Exception | None) -> None:
self.reader.feed_eof()
if exc is not None:
self.response.set_exception(exc)
self.run_parser()


async def connect_http_proxy(
proxy: Proxy,
ws_uri: WebSocketURI,
*,
user_agent_header: str | None = None,
**kwargs: Any,
) -> asyncio.Transport:
Expand All @@ -806,8 +824,8 @@ async def connect_http_proxy(
try:
# This raises exceptions if the connection to the proxy fails.
await protocol.response
except Exception:
transport.close()
except (asyncio.CancelledError, Exception):
transport.abort()
raise

return transport
Loading
Loading