File size: 5,336 Bytes
9c6594c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
from __future__ import annotations

import argparse
import asyncio
import os
import sys
from typing import Generator

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


__all__ = ["main"]


def print_during_input(string: str) -> None:
    sys.stdout.write(
        # Save cursor position
        "\N{ESC}7"
        # Add a new line
        "\N{LINE FEED}"
        # Move cursor up
        "\N{ESC}[A"
        # Insert blank line, scroll last line down
        "\N{ESC}[L"
        # Print string in the inserted blank line
        f"{string}\N{LINE FEED}"
        # Restore cursor position
        "\N{ESC}8"
        # Move cursor down
        "\N{ESC}[B"
    )
    sys.stdout.flush()


def print_over_input(string: str) -> None:
    sys.stdout.write(
        # Move cursor to beginning of line
        "\N{CARRIAGE RETURN}"
        # Delete current line
        "\N{ESC}[K"
        # Print string
        f"{string}\N{LINE FEED}"
    )
    sys.stdout.flush()


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())


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:  # pragma: no cover
            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],
            # Clean up and exit when the server closes the connection
            # or the user enters EOT (^D), whichever happens first.
            return_when=asyncio.FIRST_COMPLETED,
        )
    # asyncio.run() cancels the main task when the user triggers SIGINT (^C).
    # https://docs.python.org/3/library/asyncio-runner.html#handling-keyboard-interruption
    # Clean up and exit without re-raising CancelledError to prevent Python
    # from raising KeyboardInterrupt and displaying a stack track.
    except asyncio.CancelledError:  # pragma: no cover
        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(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(
        prog="websockets",
        description="Interactive WebSocket client.",
        add_help=False,
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--version", action="store_true")
    group.add_argument("uri", metavar="<uri>", nargs="?")
    args = parser.parse_args(argv)

    if args.version:
        print(f"websockets {websockets_version}")
        return

    if args.uri is None:
        parser.print_usage()
        sys.exit(2)

    # Enable VT100 to support ANSI escape codes in Command Prompt on Windows.
    # See https://github.com/python/cpython/issues/74261 for why this works.
    if sys.platform == "win32":
        os.system("")

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

    # Remove the try/except block when dropping Python < 3.11.
    try:
        asyncio.run(interactive_client(args.uri))
    except KeyboardInterrupt:  # pragma: no cover
        pass