- use logger, not print statement

- more document.
- add JavaScript WebSocket-like API
- add examples

5 files changed, 190 insertions, 22 deletions

diff --git a/README b/README
index f94a709..c5d48fd 100644
--- a/README
+++ b/README
@@ -8,7 +8,9 @@ Installation
 =============
 
 This module is tested on only Python 2.7.
-Type "python setup.py install" to install.
+Type "python setup.py install" or "pip install websocket-client" to install.
+
+This module does not depend on any other module.
 
 Example
 ========
@@ -22,3 +24,4 @@ Example
     print "Reeiving..."
     result =  ws.recv()
     print "Received '%s'" % result
+    ws.close()
diff --git a/examples/echo_client.py b/examples/echo_client.py
new file mode 100644
index 0000000..f8f6bf3
--- /dev/null
+++ b/examples/echo_client.py
@@ -0,0 +1,12 @@
+import websocket
+
+if __name__ == "__main__":
+    websocket.enableTrace(True)
+    ws = websocket.create_connection("ws://localhost:5000/chat")
+    print "Sending 'Hello, World'..."
+    ws.send("Hello, World")
+    print "Sent"
+    print "Receiving..."
+    result =  ws.recv()
+    print "Received '%s'" % result
+    ws.close()
diff --git a/examples/echoapp_client.py b/examples/echoapp_client.py
new file mode 100644
index 0000000..15552a3
--- /dev/null
+++ b/examples/echoapp_client.py
@@ -0,0 +1,33 @@
+import websocket
+import thread
+import time
+
+def on_message(ws, message):
+    print message
+
+def on_error(ws, error):
+    print error
+
+def on_close(ws):
+    print "### closed ###"
+
+def on_open(ws):
+    def run(*args):
+        for i in range(3):
+            time.sleep(1)
+            ws.send("Hello %d" % i)
+        time.sleep(1)
+        ws.close()
+        print "thread terminating..."
+    thread.start_new_thread(run, ())
+
+
+if __name__ == "__main__":
+    websocket.enableTrace(True)
+    ws = websocket.WebSocketApp("ws://localhost:5000/chat",
+                                on_message = on_message,
+                                on_error = on_error,
+                                on_close = on_close)
+    ws.on_open = on_open
+    
+    ws.run_forever()
diff --git a/test_websocket.py b/test_websocket.py
index 15598ed..be0f5e9 100644
--- a/test_websocket.py
+++ b/test_websocket.py
@@ -33,7 +33,7 @@ class HeaderSockMock(StringSockMock):
 
 class WebSocketTest(unittest.TestCase):
     def setUp(self):
-        pass
+        ws.enableTrace(True)
     
     def tearDown(self):
         pass
diff --git a/websocket.py b/websocket.py
index 1603e77..bb65761 100644
--- a/websocket.py
+++ b/websocket.py
@@ -3,8 +3,11 @@ from urlparse import urlparse
 import random
 import struct
 import md5
+import logging
 
 
+logger = logging.getLogger()
+
 class WebSocketException(Exception):
     pass
 
@@ -15,8 +18,15 @@ default_timeout = None
 traceEnabled = False
 
 def enableTrace(tracable):
+    """
+    turn on/off the tracability.
+    """
     global traceEnabled
     traceEnabled = tracable
+    if tracable:
+        if not logger.handlers:
+            logger.addHandler(logging.StreamHandler())
+        logger.setLevel(logging.DEBUG)
 
 def setdefaulttimeout(timeout):
     """
@@ -119,7 +129,7 @@ HEADERS_TO_EXIST_FOR_HIXIE75 = [
     "websocket-location",
 ]
 
-class SSLSocketWrapper(object):
+class _SSLSocketWrapper(object):
     def __init__(self, sock):
         self.ssl = socket.ssl(sock)
 
@@ -130,6 +140,23 @@ class SSLSocketWrapper(object):
         return self.ssl.write(payload)
 
 class WebSocket(object):
+    """
+    Low level WebSocket interface.
+    This class is based on
+      The WebSocket protocol draft-hixie-thewebsocketprotocol-76
+      http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
+
+    We can connect to the websocket server and send/recieve data.
+    The following example is a echo client.
+
+    >>> import websocket
+    >>> ws = websocket.WebSocket()
+    >>> ws.Connect("ws://localhost:8080/echo")
+    >>> ws.send("Hello, Server")
+    >>> ws.recv()
+    'Hello, Server'
+    >>> ws.close()
+    """
     def __init__(self):
         """
         Initalize WebSocket object.
@@ -157,7 +184,7 @@ class WebSocket(object):
         # TODO: we need to support proxy
         self.sock.connect((hostname, port))
         if is_secure:
-            self.io_sock = SSLSocketWrapper(self.sock)
+            self.io_sock = _SSLSocketWrapper(self.sock)
         self._handshake(hostname, port, resource, **options)
 
     def _handshake(self, host, port, resource, **options):
@@ -187,9 +214,9 @@ class WebSocket(object):
         header_str = "\r\n".join(headers)
         sock.send(header_str)
         if traceEnabled:
-            print "--- request header ---"
-            print header_str
-            print "-----------------------"
+            logger.debug( "--- request header ---")
+            logger.debug( header_str)
+            logger.debug("-----------------------")
 
         status, resp_headers = self._read_headers()
         if status != 101:
@@ -209,8 +236,8 @@ class WebSocket(object):
         self.connected = True
     
     def _validate_resp(self, number_1, number_2, key3, resp):
-        challenge = struct.pack("!i", number_1)
-        challenge += struct.pack("!i", number_2)
+        challenge = struct.pack("!I", number_1)
+        challenge += struct.pack("!I", number_2)
         challenge += key3
         digest = md5.md5(challenge).digest()
         
@@ -219,9 +246,9 @@ class WebSocket(object):
     def _get_resp(self):
         result = self._recv(16)
         if traceEnabled:
-            print "--- challenge response result ---"
-            print repr(result)
-            print "---------------------------------"
+            logger.debug("--- challenge response result ---")
+            logger.debug(repr(result))
+            logger.debug("---------------------------------")
         
         return result
 
@@ -255,7 +282,7 @@ class WebSocket(object):
         status = None
         headers = {}
         if traceEnabled:
-            print "--- response header ---"
+            logger.debug("--- response header ---")
             
         while True:
             line = self._recv_line()
@@ -263,7 +290,7 @@ class WebSocket(object):
                 break
             line = line.strip()
             if traceEnabled:
-                print line
+                logger.debug(line)
             if not status:
                 status_info = line.split(" ", 2)
                 status = int(status_info[1])
@@ -276,7 +303,7 @@ class WebSocket(object):
                     raise WebSocketException("Invalid header")
 
         if traceEnabled:
-            print "-----------------------"
+            logger.debug("-----------------------")
         
         return status, headers    
     
@@ -286,13 +313,18 @@ class WebSocket(object):
         """
         if isinstance(payload, unicode):
             payload = payload.encode("utf-8")
-        self.io_sock.send("".join(["\x00", payload, "\xff"]))
+        data = "".join(["\x00", payload, "\xff"])
+        self.io_sock.send(data)
+        if traceEnabled:
+            logger.debug("send: " + repr(data))
 
     def recv(self):
         """
         Reeive utf-8 string data from the server.
         """
         b = self._recv(1)
+        if enableTrace:
+            logger.debug("recv frame: " + repr(b))
         frame_type = ord(b)
         if frame_type == 0x00:
             bytes = []
@@ -303,11 +335,15 @@ class WebSocket(object):
                 else:
                     bytes.append(b)
             return "".join(bytes)
-        elif frame_type > 0x80:
+        elif 0x80 < frame_type < 0xff:
             # which frame type is valid?
             length = self._read_length()
             bytes = self._recv_strict(length)
             return bytes
+        elif frame_type == 0xff:
+            n = self._recv(1)
+            self._closeInternal()
+            return None
         else:
             raise WebSocketException("Invalid frame type")
 
@@ -328,11 +364,22 @@ class WebSocket(object):
         if self.connected:
             try:
                 self.io_sock.send("\xff\x00")
-                result = self._recv(2)
-                if result != "\xff\x00":
-                    logger.error("bad closing Handshake")
+                timeout = self.sock.gettimeout()
+                self.sock.settimeout(1)
+                try:
+                    result = self._recv(2)
+                    if result != "\xff\x00":
+                        logger.error("bad closing Handshake")
+                except:
+                    pass
+                self.sock.settimeout(timeout)
+                self.sock.shutdown(socket.SHUT_RDWR)
             except:
                 pass
+        self._closeInternal()
+
+    def _closeInternal(self):
+        self.connected = False
         self.sock.close()
         self.io_sock = self.sock
         
@@ -360,11 +407,83 @@ class WebSocket(object):
                 break
         return "".join(line)
             
-        
+class WebSocketApp(object):
+    """
+    Higher level of APIs are provided. 
+    The interface is like JavaScript WebSocket object.
+    """
+    def __init__(self, url,
+                 on_open = None, on_message = None, on_error = None, 
+                 on_close = None):
+        """
+        url: websocket url.
+        on_open: callable object which is called at opening websocket.
+          this function has one argument. The arugment is this class object.
+        on_message: callbale object which is called when recieved data.
+         on_message has 2 arguments. 
+         The 1st arugment is this class object.
+         The passing 2nd arugment is utf-8 string which we get from the server.
+       on_error: callable object which is called when we get error.
+         on_error has 2 arguments.
+         The 1st arugment is this class object.
+         The passing 2nd arugment is exception object.
+       on_close: callable object which is called when closed the connection.
+         this function has one argument. The arugment is this class object.
+        """
+        self.url = url
+        self.on_open = on_open
+        self.on_message = on_message
+        self.on_error = on_error
+        self.on_close = on_close
+        self.sock = None
+
+    def send(self, data):
+        """
+        send message. data must be utf-8 string or unicode.
+        """
+        self.sock.send(data)
+
+    def close(self):
+        """
+        close websocket connection.
+        """
+        self.sock.close()
+
+    def run_forever(self):
+        """
+        run event loop for WebSocket framework.
+        This loop is infinite loop and is alive during websocket is available.
+        """
+        if self.sock:
+            raise WebSocketException("socket is already opened")
+        try:
+            self.sock = WebSocket()
+            self.sock.connect(self.url)
+            self._run_with_no_err(self.on_open)
+            while True:
+                data = self.sock.recv()
+                if data is None:
+                    break
+                self._run_with_no_err(self.on_message, data)
+        except Exception, e:
+            self._run_with_no_err(self.on_error, e)
+        finally:
+            self.sock.close()
+            self._run_with_no_err(self.on_close)
+            self.sock = None
+
+    def _run_with_no_err(self, callback, *args):
+        if callback:
+            try:
+                callback(self, *args)
+            except Exception, e:
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.error(e)
+
 
 if __name__ == "__main__":
     enableTrace(True)
-    # ws = create_connection("ws://localhost:8080/echo")
+    #ws = create_connection("ws://localhost:8080/echo")
     ws = create_connection("ws://localhost:5000/chat")
     print "Sending 'Hello, World'..."
     ws.send("Hello, World")
@@ -372,6 +491,7 @@ if __name__ == "__main__":
     print "Receiving..."
     result =  ws.recv()
     print "Received '%s'" % result
+    ws.close()