Writing non blocking command line client

[kracekumar]

I am trying to create command line client for websocket message. User can type the message and mean if any message is received it should be printed to terminal.

Here is a example using gevent.

# -*- coding: utf-8 -*-
# cli.py

import sys
import os
import fcntl

import json
import urllib
import urlparse

import gevent
from gevent.socket import wait_read
from gevent.event import Event

import select

from ws4py.client.geventclient import WebSocketClient
from ws4py.exc import HandshakeError


class StdioPipedWebSocketHelper:
    def received_message(self, websocket, m):
        if m.is_binary:
            print("Binary data receieved. Skipping")
        else:
            sys.stdout.write(m.data)
        sys.stdout.write("\n")
        sys.stdout.flush()

    def opened(self, websocket):
        def connect_stdin():
            sys.stdout.write("\n")
            fcntl.fcntl(sys.stdin, fcntl.F_SETFL, os.O_NONBLOCK)
            while True:
                wait_read(sys.stdin.fileno())
                print(">")
                buf = sys.stdin.readline()
                if len(buf) == 0:
                    break
                print("[sending payload of length %d]" % (len(buf)))
                websocket.send_msg(buf)

        # XXX: We wait for the socket to open before reading stdin so that we
        # support behaviour like: echo foo | wssh -l ...
        gevent.spawn(connect_stdin)
        print("done")

    def closed(self):
        print("Connection closed by peer.")
        print("Shutting down...")
        exit(0)


class Client(WebSocketClient):
    def __init__(self, credentials):
        url = credentials.pop('url')
        qs = urllib.urlencode(credentials)
        full_url = url + "?" + qs
        print("Connecting: {}".format(full_url))
        WebSocketClient.__init__(self, full_url)

        self.shutdown_cond = Event()
        self.io = StdioPipedWebSocketHelper()

    def opened(self):
        self.io.opened(self)

    def closed(self, closed, reason):
        self.io.closed()
        self.shutdown_cond.set()

    def received_message(self, message):
        msg_obj = parse_msg(message)
        self.io.received_message(message=msg_obj)

    def send_msg(message):
        msg = json.dumps({'post': message})
        self.send(msg, False)

    def connect_and_wait(self):
        self.connect()
        self.shutdown_cond.wait()


def start(credentials):
    """Connect to the serve and keep running.

    :param credentials: `dict` containing details about server and tokens.
    """
    client = Client(credentials=credentials)
    try:
        client.connect_and_wait()
    except (IOError, HandshakeError), e:
        print >> sys.stderr, e



if __name__ == "__main__":
    credentials = {'url': 'ws://echo.websocket.org', 'foo': 'bar'}
    start(credentials=credentials)

Most of the code is borrowed from @progrium's work on wssh. wssh repo no more work because of old version of ws4py.

(ws4py_tests)➜  ws4py_test  python cli.py
Connecting: ws://echo.websocket.org?foo=bar
done
foo
foo
bar
^CConnection closed by peer.
Shutting down...

After connecting to server, code is blocked inside websocket.once. Input from terminal is not read and transmitted.

Trace

(ws4py_tests)➜  ws4py_test  python -m trace --trace cli.py
...
 --- modulename: greenlet, funcname: start
greenlet.py(186):         if self._start_event is None:
greenlet.py(187):             self._start_event = self.parent.loop.run_callback(self.switch)
greenlet.py(203):         return g
cli.py(46):         print("done")
done
websocket.py(426):                 while not self.terminated:
 --- modulename: websocket, funcname: terminated
websocket.py(194):         return self.client_terminated is True and self.server_terminated is True
websocket.py(427):                     if not self.once():
 --- modulename: websocket, funcname: once
websocket.py(295):         if self.terminated:
 --- modulename: websocket, funcname: terminated
websocket.py(194):         return self.client_terminated is True and self.server_terminated is True
websocket.py(299):         try:
websocket.py(300):             b = self.sock.recv(self.reading_buffer_size)

Venv

(ws4py_tests)➜  ws4py_test  pip freeze
gevent==1.0.1
greenlet==0.4.7
ws4py==0.3.4

Any thing I am doing dumb here ?