__main__.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
  1. import argparse
  2. import asyncio
  3. import os
  4. import signal
  5. import sys
  6. import threading
  7. from typing import Any, Set
  8. from .exceptions import ConnectionClosed, format_close
  9. from .legacy.client import connect
  10. if sys.platform == "win32":
  11. def win_enable_vt100() -> None:
  12. """
  13. Enable VT-100 for console output on Windows.
  14. See also https://bugs.python.org/issue29059.
  15. """
  16. import ctypes
  17. STD_OUTPUT_HANDLE = ctypes.c_uint(-11)
  18. INVALID_HANDLE_VALUE = ctypes.c_uint(-1)
  19. ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x004
  20. handle = ctypes.windll.kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
  21. if handle == INVALID_HANDLE_VALUE:
  22. raise RuntimeError("unable to obtain stdout handle")
  23. cur_mode = ctypes.c_uint()
  24. if ctypes.windll.kernel32.GetConsoleMode(handle, ctypes.byref(cur_mode)) == 0:
  25. raise RuntimeError("unable to query current console mode")
  26. # ctypes ints lack support for the required bit-OR operation.
  27. # Temporarily convert to Py int, do the OR and convert back.
  28. py_int_mode = int.from_bytes(cur_mode, sys.byteorder)
  29. new_mode = ctypes.c_uint(py_int_mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING)
  30. if ctypes.windll.kernel32.SetConsoleMode(handle, new_mode) == 0:
  31. raise RuntimeError("unable to set console mode")
  32. def exit_from_event_loop_thread(
  33. loop: asyncio.AbstractEventLoop, stop: "asyncio.Future[None]"
  34. ) -> None:
  35. loop.stop()
  36. if not stop.done():
  37. # When exiting the thread that runs the event loop, raise
  38. # KeyboardInterrupt in the main thread to exit the program.
  39. if sys.platform == "win32":
  40. ctrl_c = signal.CTRL_C_EVENT
  41. else:
  42. ctrl_c = signal.SIGINT
  43. os.kill(os.getpid(), ctrl_c)
  44. def print_during_input(string: str) -> None:
  45. sys.stdout.write(
  46. # Save cursor position
  47. "\N{ESC}7"
  48. # Add a new line
  49. "\N{LINE FEED}"
  50. # Move cursor up
  51. "\N{ESC}[A"
  52. # Insert blank line, scroll last line down
  53. "\N{ESC}[L"
  54. # Print string in the inserted blank line
  55. f"{string}\N{LINE FEED}"
  56. # Restore cursor position
  57. "\N{ESC}8"
  58. # Move cursor down
  59. "\N{ESC}[B"
  60. )
  61. sys.stdout.flush()
  62. def print_over_input(string: str) -> None:
  63. sys.stdout.write(
  64. # Move cursor to beginning of line
  65. "\N{CARRIAGE RETURN}"
  66. # Delete current line
  67. "\N{ESC}[K"
  68. # Print string
  69. f"{string}\N{LINE FEED}"
  70. )
  71. sys.stdout.flush()
  72. async def run_client(
  73. uri: str,
  74. loop: asyncio.AbstractEventLoop,
  75. inputs: "asyncio.Queue[str]",
  76. stop: "asyncio.Future[None]",
  77. ) -> None:
  78. try:
  79. websocket = await connect(uri)
  80. except Exception as exc:
  81. print_over_input(f"Failed to connect to {uri}: {exc}.")
  82. exit_from_event_loop_thread(loop, stop)
  83. return
  84. else:
  85. print_during_input(f"Connected to {uri}.")
  86. try:
  87. while True:
  88. incoming: asyncio.Future[Any] = asyncio.ensure_future(websocket.recv())
  89. outgoing: asyncio.Future[Any] = asyncio.ensure_future(inputs.get())
  90. done: Set[asyncio.Future[Any]]
  91. pending: Set[asyncio.Future[Any]]
  92. done, pending = await asyncio.wait(
  93. [incoming, outgoing, stop], return_when=asyncio.FIRST_COMPLETED
  94. )
  95. # Cancel pending tasks to avoid leaking them.
  96. if incoming in pending:
  97. incoming.cancel()
  98. if outgoing in pending:
  99. outgoing.cancel()
  100. if incoming in done:
  101. try:
  102. message = incoming.result()
  103. except ConnectionClosed:
  104. break
  105. else:
  106. if isinstance(message, str):
  107. print_during_input("< " + message)
  108. else:
  109. print_during_input("< (binary) " + message.hex())
  110. if outgoing in done:
  111. message = outgoing.result()
  112. await websocket.send(message)
  113. if stop in done:
  114. break
  115. finally:
  116. await websocket.close()
  117. close_status = format_close(websocket.close_code, websocket.close_reason)
  118. print_over_input(f"Connection closed: {close_status}.")
  119. exit_from_event_loop_thread(loop, stop)
  120. def main() -> None:
  121. # If we're on Windows, enable VT100 terminal support.
  122. if sys.platform == "win32":
  123. try:
  124. win_enable_vt100()
  125. except RuntimeError as exc:
  126. sys.stderr.write(
  127. f"Unable to set terminal to VT100 mode. This is only "
  128. f"supported since Win10 anniversary update. Expect "
  129. f"weird symbols on the terminal.\nError: {exc}\n"
  130. )
  131. sys.stderr.flush()
  132. try:
  133. import readline # noqa
  134. except ImportError: # Windows has no `readline` normally
  135. pass
  136. # Parse command line arguments.
  137. parser = argparse.ArgumentParser(
  138. prog="python -m websockets",
  139. description="Interactive WebSocket client.",
  140. add_help=False,
  141. )
  142. parser.add_argument("uri", metavar="<uri>")
  143. args = parser.parse_args()
  144. # Create an event loop that will run in a background thread.
  145. loop = asyncio.new_event_loop()
  146. # Due to zealous removal of the loop parameter in the Queue constructor,
  147. # we need a factory coroutine to run in the freshly created event loop.
  148. async def queue_factory() -> "asyncio.Queue[str]":
  149. return asyncio.Queue()
  150. # Create a queue of user inputs. There's no need to limit its size.
  151. inputs: "asyncio.Queue[str]" = loop.run_until_complete(queue_factory())
  152. # Create a stop condition when receiving SIGINT or SIGTERM.
  153. stop: asyncio.Future[None] = loop.create_future()
  154. # Schedule the task that will manage the connection.
  155. asyncio.ensure_future(run_client(args.uri, loop, inputs, stop), loop=loop)
  156. # Start the event loop in a background thread.
  157. thread = threading.Thread(target=loop.run_forever)
  158. thread.start()
  159. # Read from stdin in the main thread in order to receive signals.
  160. try:
  161. while True:
  162. # Since there's no size limit, put_nowait is identical to put.
  163. message = input("> ")
  164. loop.call_soon_threadsafe(inputs.put_nowait, message)
  165. except (KeyboardInterrupt, EOFError): # ^C, ^D
  166. loop.call_soon_threadsafe(stop.set_result, None)
  167. # Wait for the event loop to terminate.
  168. thread.join()
  169. # For reasons unclear, even though the loop is closed in the thread,
  170. # it still thinks it's running here.
  171. loop.close()
  172. if __name__ == "__main__":
  173. main()