Implemented connection hooks, cleaned up some code

connection.py

@@ -49,6 +49,9 @@ class Connection(object):
         self.ping_sent = False
         self.ping_payload = None
 
+        self.hooks_send = []
+        self.hooks_recv = []
+
         self.onopen()
 
     def send(self, message, fragment_size=None, mask=False):
@@ -57,6 +60,9 @@ class Connection(object):
         fragmented into multiple frames whose payload size does not extend
         `fragment_size`.
         """
+        for hook in self.hooks_send:
+            message = hook(message)
+
         if fragment_size is None:
             self.sock.send(message.frame(mask=mask))
         else:
@@ -87,7 +93,12 @@ class Connection(object):
         for f in fragments:
             payload += f.payload
 
-        return create_message(fragments[0].opcode, payload)
+        message = create_message(fragments[0].opcode, payload)
+
+        for hook in self.hooks_recv:
+            message = hook(message)
+
+        return message
 
     def handle_control_frame(self, frame):
         """
@@ -191,6 +202,36 @@ class Connection(object):
         self.onclose(code, reason)
         self.sock.close()
 
+    def add_hook(self, send=None, recv=None, prepend=False):
+        """
+        Add a pair of send and receive hooks that are called for each frame
+        that is sent or received. A hook is a function that receives a single
+        argument - a Message instance - and returns a `Message` instance as
+        well.
+
+        `prepend` is a flag indicating whether the send hook is prepended to
+        the other send hooks.
+
+        For example, to add an automatic JSON conversion to messages and
+        eliminate the need to contruct TextMessage instances to all messages:
+        >>> import twspy, json
+        >>> conn = Connection(...)
+        >>> conn.add_hook(lambda data: tswpy.TextMessage(json.dumps(data)),
+        >>>               lambda message: json.loads(message.payload))
+        >>> conn.send({'foo': 'bar'})  # Sends text message {"foo":"bar"}
+        >>> conn.recv()                # May be dict(foo='bar')
+
+        Note that here `prepend=True`, so that data passed to `send()` is first
+        encoded and then packed into a frame. Of course, one could also decide
+        to add the base64 hook first, or to return a new `Frame` instance with
+        base64-encoded data.
+        """
+        if send:
+            self.hooks_send.insert(0 if prepend else -1, send)
+
+        if recv:
+            self.hooks_recv.insert(-1 if prepend else 0, recv)
+
     def onopen(self):
         """
         Called after the connection is initialized.

websocket.py

@@ -65,12 +65,16 @@ class websocket(object):
         self.trusted_origins = trusted_origins
         self.location = location
         self.auth = auth
-        self.sock = sock or socket.socket(sfamily, socket.SOCK_STREAM, sproto)
+
         self.secure = False
+
         self.handshake_sent = False
+
         self.hooks_send = []
         self.hooks_recv = []
 
+        self.sock = sock or socket.socket(sfamily, socket.SOCK_STREAM, sproto)
+
     def bind(self, address):
         self.sock.bind(address)
 
@@ -165,14 +169,17 @@ class websocket(object):
         that is sent or received. A hook is a function that receives a single
         argument - a Frame instance - and returns a `Frame` instance as well.
 
-        `prepend` is a flag indicating whether the `send` hook
+        `prepend` is a flag indicating whether the send hook is prepended to
+        the other send hooks. This is expecially useful when a program uses
+        extensions such as the built-in `DeflateFrame` extension. These
+        extensions are installed using these hooks as well.
 
         For example, the following code creates a `Frame` instance for data
         being sent and removes the instance for received data. This way, data
         can be sent and received as if on a regular socket.
         >>> import twspy
         >>> sock.add_hook(lambda data: tswpy.Frame(tswpy.OPCODE_TEXT, data),
-        >>>               lambda f: f.data)
+        >>>               lambda frame: frame.payload)
 
         To add base64 encoding to the example above:
         >>> import base64