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
|