GitProtogen
109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
import asyncio
|
|
import threading
|
|
import websockets
|
|
from prompt_toolkit import PromptSession
|
|
from prompt_toolkit.key_binding import KeyBindings
|
|
from prompt_toolkit.keys import Keys
|
|
|
|
URI = "ws://localhost:8080/ws"
|
|
RETRY_DELAY = 2
|
|
|
|
bindings = KeyBindings()
|
|
|
|
@bindings.add(Keys.Enter)
|
|
def submit(event):
|
|
event.current_buffer.validate_and_handle()
|
|
|
|
@bindings.add(Keys.Escape, Keys.Enter)
|
|
def newline(event):
|
|
event.current_buffer.insert_text("\n")
|
|
|
|
send_queue: asyncio.Queue = None
|
|
loop: asyncio.AbstractEventLoop = None
|
|
shutdown_event: asyncio.Event = None
|
|
|
|
async def receiver(ws):
|
|
async for message in ws:
|
|
print(f"\n[SERVER] {message}", flush=True)
|
|
|
|
async def sender(ws):
|
|
while True:
|
|
msg = await send_queue.get()
|
|
if msg is None:
|
|
break
|
|
await ws.send(msg)
|
|
|
|
async def run():
|
|
global send_queue
|
|
send_queue = asyncio.Queue()
|
|
|
|
input_thread = threading.Thread(target=input_loop, daemon=True)
|
|
input_thread_started = False
|
|
|
|
while not shutdown_event.is_set():
|
|
try:
|
|
async with websockets.connect(URI) as ws:
|
|
print(f"Connected to {URI}")
|
|
if not input_thread_started:
|
|
print("Alt+Enter = newline | Enter = send | Ctrl+C = quit\n")
|
|
input_thread.start()
|
|
input_thread_started = True
|
|
|
|
recv_task = asyncio.create_task(receiver(ws))
|
|
send_task = asyncio.create_task(sender(ws))
|
|
shutdown_task = asyncio.create_task(shutdown_event.wait())
|
|
|
|
done, pending = await asyncio.wait(
|
|
[recv_task, send_task, shutdown_task],
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
|
|
for t in pending:
|
|
t.cancel()
|
|
|
|
if shutdown_event.is_set():
|
|
return
|
|
|
|
# reconnect
|
|
send_queue = asyncio.Queue()
|
|
print(f"\n[Disconnected] Reconnecting in {RETRY_DELAY}s...")
|
|
|
|
except (OSError, websockets.exceptions.WebSocketException):
|
|
if shutdown_event.is_set():
|
|
return
|
|
print(f"[Waiting for server] Retrying in {RETRY_DELAY}s...", flush=True)
|
|
|
|
try:
|
|
await asyncio.wait_for(shutdown_event.wait(), timeout=RETRY_DELAY)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
def input_loop():
|
|
session = PromptSession(
|
|
key_bindings=bindings,
|
|
multiline=True,
|
|
)
|
|
while True:
|
|
try:
|
|
text = session.prompt(">>> ")
|
|
if text is not None:
|
|
asyncio.run_coroutine_threadsafe(send_queue.put(text), loop)
|
|
except (EOFError, KeyboardInterrupt):
|
|
asyncio.run_coroutine_threadsafe(shutdown_event.set(), loop)
|
|
break
|
|
|
|
def main():
|
|
global loop, shutdown_event
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
shutdown_event = asyncio.Event()
|
|
try:
|
|
loop.run_until_complete(run())
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|