Skip to content

Commit

Permalink
Rewrite interactive client (again) without threads.
Browse files Browse the repository at this point in the history
Fix #1592.
  • Loading branch information
aaugustin committed Feb 16, 2025
1 parent bc3fd29 commit a1ba01d
Showing 1 changed file with 91 additions and 47 deletions.
138 changes: 91 additions & 47 deletions src/websockets/__main__.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,16 @@
from __future__ import annotations

import argparse
import asyncio
import os
import signal
import sys
import threading


try:
import readline # noqa: F401
except ImportError: # readline isn't available on all platforms
pass
from typing import Generator

from .asyncio.client import ClientConnection, connect
from .asyncio.messages import SimpleQueue
from .exceptions import ConnectionClosed
from .frames import Close
from .sync.client import ClientConnection, connect
from .streams import StreamReader
from .version import version as websockets_version


Expand Down Expand Up @@ -49,24 +46,94 @@ def print_over_input(string: str) -> None:
sys.stdout.flush()


def print_incoming_messages(websocket: ClientConnection, stop: threading.Event) -> None:
for message in websocket:
class ReadLines(asyncio.Protocol):
def __init__(self) -> None:
self.reader = StreamReader()
self.messages: SimpleQueue[str] = SimpleQueue()

def parse(self) -> Generator[None, None, None]:
while True:
sys.stdout.write("> ")
sys.stdout.flush()
line = yield from self.reader.read_line(sys.maxsize)
self.messages.put(line.decode().rstrip("\r\n"))

def connection_made(self, transport: asyncio.BaseTransport) -> None:
self.parser = self.parse()
next(self.parser)

def data_received(self, data: bytes) -> None:
self.reader.feed_data(data)
next(self.parser)

def eof_received(self) -> None:
self.reader.feed_eof()
# next(self.parser) isn't useful and would raise EOFError.

def connection_lost(self, exc: Exception | None) -> None:
self.reader.discard()
self.messages.abort()


async def print_incoming_messages(websocket: ClientConnection) -> None:
async for message in websocket:
if isinstance(message, str):
print_during_input("< " + message)
else:
print_during_input("< (binary) " + message.hex())
if not stop.is_set():
# When the server closes the connection, raise KeyboardInterrupt
# in the main thread to exit the program.
if sys.platform == "win32":
ctrl_c = signal.CTRL_C_EVENT
else:
ctrl_c = signal.SIGINT
os.kill(os.getpid(), ctrl_c)


async def send_outgoing_messages(
websocket: ClientConnection,
messages: SimpleQueue[str],
) -> None:
while True:
try:
message = await messages.get()
except EOFError:
break
try:
await websocket.send(message)
except ConnectionClosed:
break


async def interactive_client(uri: str) -> None:
try:
websocket = await connect(uri)
except Exception as exc:
print(f"Failed to connect to {uri}: {exc}.")
sys.exit(1)
else:
print(f"Connected to {uri}.")

loop = asyncio.get_running_loop()
transport, protocol = await loop.connect_read_pipe(ReadLines, sys.stdin)
incoming = asyncio.create_task(
print_incoming_messages(websocket),
)
outgoing = asyncio.create_task(
send_outgoing_messages(websocket, protocol.messages),
)
try:
await asyncio.wait(
[incoming, outgoing],
return_when=asyncio.FIRST_COMPLETED,
)
except (KeyboardInterrupt, EOFError): # ^C, ^D
pass
finally:
incoming.cancel()
outgoing.cancel()
transport.close()

await websocket.close()
assert websocket.close_code is not None and websocket.close_reason is not None
close_status = Close(websocket.close_code, websocket.close_reason)
print_over_input(f"Connection closed: {close_status}.")


def main() -> None:
# Parse command line arguments.
parser = argparse.ArgumentParser(
prog="python -m websockets",
description="Interactive WebSocket client.",
Expand All @@ -90,34 +157,11 @@ def main() -> None:
os.system("")

try:
websocket = connect(args.uri)
except Exception as exc:
print(f"Failed to connect to {args.uri}: {exc}.")
sys.exit(1)
else:
print(f"Connected to {args.uri}.")

stop = threading.Event()

# Start the thread that reads messages from the connection.
thread = threading.Thread(target=print_incoming_messages, args=(websocket, stop))
thread.start()

# Read from stdin in the main thread in order to receive signals.
try:
while True:
# Since there's no size limit, put_nowait is identical to put.
message = input("> ")
websocket.send(message)
except (KeyboardInterrupt, EOFError): # ^C, ^D
stop.set()
websocket.close()

assert websocket.close_code is not None and websocket.close_reason is not None
close_status = Close(websocket.close_code, websocket.close_reason)
print_over_input(f"Connection closed: {close_status}.")
import readline # noqa: F401
except ImportError: # readline isn't available on all platforms
pass

thread.join()
asyncio.run(interactive_client(args.uri))


if __name__ == "__main__":
Expand Down

0 comments on commit a1ba01d

Please sign in to comment.