Rewrite interactive client (again) without threads.

Fix #1592.

Author: aaugustin

src/websockets/__main__.py

@@ -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
 
 
@@ -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.",
@@ -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__":