|
1 |
| -from __future__ import annotations |
2 |
| - |
3 |
| -import argparse |
4 |
| -import asyncio |
5 |
| -import os |
6 |
| -import sys |
7 |
| -from typing import Generator |
8 |
| - |
9 |
| -from .asyncio.client import ClientConnection, connect |
10 |
| -from .asyncio.messages import SimpleQueue |
11 |
| -from .exceptions import ConnectionClosed |
12 |
| -from .frames import Close |
13 |
| -from .streams import StreamReader |
14 |
| -from .version import version as websockets_version |
15 |
| - |
16 |
| - |
17 |
| -def print_during_input(string: str) -> None: |
18 |
| - sys.stdout.write( |
19 |
| - # Save cursor position |
20 |
| - "\N{ESC}7" |
21 |
| - # Add a new line |
22 |
| - "\N{LINE FEED}" |
23 |
| - # Move cursor up |
24 |
| - "\N{ESC}[A" |
25 |
| - # Insert blank line, scroll last line down |
26 |
| - "\N{ESC}[L" |
27 |
| - # Print string in the inserted blank line |
28 |
| - f"{string}\N{LINE FEED}" |
29 |
| - # Restore cursor position |
30 |
| - "\N{ESC}8" |
31 |
| - # Move cursor down |
32 |
| - "\N{ESC}[B" |
33 |
| - ) |
34 |
| - sys.stdout.flush() |
35 |
| - |
36 |
| - |
37 |
| -def print_over_input(string: str) -> None: |
38 |
| - sys.stdout.write( |
39 |
| - # Move cursor to beginning of line |
40 |
| - "\N{CARRIAGE RETURN}" |
41 |
| - # Delete current line |
42 |
| - "\N{ESC}[K" |
43 |
| - # Print string |
44 |
| - f"{string}\N{LINE FEED}" |
45 |
| - ) |
46 |
| - sys.stdout.flush() |
47 |
| - |
48 |
| - |
49 |
| -class ReadLines(asyncio.Protocol): |
50 |
| - def __init__(self) -> None: |
51 |
| - self.reader = StreamReader() |
52 |
| - self.messages: SimpleQueue[str] = SimpleQueue() |
53 |
| - |
54 |
| - def parse(self) -> Generator[None, None, None]: |
55 |
| - while True: |
56 |
| - sys.stdout.write("> ") |
57 |
| - sys.stdout.flush() |
58 |
| - line = yield from self.reader.read_line(sys.maxsize) |
59 |
| - self.messages.put(line.decode().rstrip("\r\n")) |
60 |
| - |
61 |
| - def connection_made(self, transport: asyncio.BaseTransport) -> None: |
62 |
| - self.parser = self.parse() |
63 |
| - next(self.parser) |
64 |
| - |
65 |
| - def data_received(self, data: bytes) -> None: |
66 |
| - self.reader.feed_data(data) |
67 |
| - next(self.parser) |
68 |
| - |
69 |
| - def eof_received(self) -> None: |
70 |
| - self.reader.feed_eof() |
71 |
| - # next(self.parser) isn't useful and would raise EOFError. |
72 |
| - |
73 |
| - def connection_lost(self, exc: Exception | None) -> None: |
74 |
| - self.reader.discard() |
75 |
| - self.messages.abort() |
76 |
| - |
77 |
| - |
78 |
| -async def print_incoming_messages(websocket: ClientConnection) -> None: |
79 |
| - async for message in websocket: |
80 |
| - if isinstance(message, str): |
81 |
| - print_during_input("< " + message) |
82 |
| - else: |
83 |
| - print_during_input("< (binary) " + message.hex()) |
84 |
| - |
85 |
| - |
86 |
| -async def send_outgoing_messages( |
87 |
| - websocket: ClientConnection, |
88 |
| - messages: SimpleQueue[str], |
89 |
| -) -> None: |
90 |
| - while True: |
91 |
| - try: |
92 |
| - message = await messages.get() |
93 |
| - except EOFError: |
94 |
| - break |
95 |
| - try: |
96 |
| - await websocket.send(message) |
97 |
| - except ConnectionClosed: |
98 |
| - break |
99 |
| - |
100 |
| - |
101 |
| -async def interactive_client(uri: str) -> None: |
102 |
| - try: |
103 |
| - websocket = await connect(uri) |
104 |
| - except Exception as exc: |
105 |
| - print(f"Failed to connect to {uri}: {exc}.") |
106 |
| - sys.exit(1) |
107 |
| - else: |
108 |
| - print(f"Connected to {uri}.") |
109 |
| - |
110 |
| - loop = asyncio.get_running_loop() |
111 |
| - transport, protocol = await loop.connect_read_pipe(ReadLines, sys.stdin) |
112 |
| - incoming = asyncio.create_task( |
113 |
| - print_incoming_messages(websocket), |
114 |
| - ) |
115 |
| - outgoing = asyncio.create_task( |
116 |
| - send_outgoing_messages(websocket, protocol.messages), |
117 |
| - ) |
118 |
| - try: |
119 |
| - await asyncio.wait( |
120 |
| - [incoming, outgoing], |
121 |
| - return_when=asyncio.FIRST_COMPLETED, |
122 |
| - ) |
123 |
| - except (KeyboardInterrupt, EOFError): # ^C, ^D |
124 |
| - pass |
125 |
| - finally: |
126 |
| - incoming.cancel() |
127 |
| - outgoing.cancel() |
128 |
| - transport.close() |
129 |
| - |
130 |
| - await websocket.close() |
131 |
| - assert websocket.close_code is not None and websocket.close_reason is not None |
132 |
| - close_status = Close(websocket.close_code, websocket.close_reason) |
133 |
| - print_over_input(f"Connection closed: {close_status}.") |
134 |
| - |
135 |
| - |
136 |
| -def main() -> None: |
137 |
| - parser = argparse.ArgumentParser( |
138 |
| - prog="websockets", |
139 |
| - description="Interactive WebSocket client.", |
140 |
| - add_help=False, |
141 |
| - ) |
142 |
| - group = parser.add_mutually_exclusive_group() |
143 |
| - group.add_argument("--version", action="store_true") |
144 |
| - group.add_argument("uri", metavar="<uri>", nargs="?") |
145 |
| - args = parser.parse_args() |
146 |
| - |
147 |
| - if args.version: |
148 |
| - print(f"websockets {websockets_version}") |
149 |
| - return |
150 |
| - |
151 |
| - if args.uri is None: |
152 |
| - parser.error("the following arguments are required: <uri>") |
153 |
| - |
154 |
| - # Enable VT100 to support ANSI escape codes in Command Prompt on Windows. |
155 |
| - # See https://github.com/python/cpython/issues/74261 for why this works. |
156 |
| - if sys.platform == "win32": # pragma: no cover |
157 |
| - os.system("") |
158 |
| - |
159 |
| - try: |
160 |
| - import readline # noqa: F401 |
161 |
| - except ImportError: # readline isn't available on all platforms |
162 |
| - pass |
163 |
| - |
164 |
| - asyncio.run(interactive_client(args.uri)) |
| 1 | +from .cli import main |
165 | 2 |
|
166 | 3 |
|
167 | 4 | if __name__ == "__main__":
|
|
0 commit comments