Add pack/unpack encoding arguments and more Python3 support. Thanks to followers in Twitter!

Author: repeatedly

msgpackrpc/client.py

@@ -7,9 +7,9 @@ class Client(session.Session):
     Client is usaful for MessagePack RPC API.
     """
 
-    def __init__(self, address, timeout=10, loop=None, builder=tcp):
+    def __init__(self, address, timeout=10, loop=None, builder=tcp, reconnect_limit=5, pack_encoding='utf-8', unpack_encoding=None):
         loop = loop or Loop()
-        session.Session.__init__(self, address, timeout, loop, builder)
+        session.Session.__init__(self, address, timeout, loop, builder, reconnect_limit, pack_encoding, unpack_encoding)
 
         if timeout:
             loop.attach_periodic_callback(self.step_timeout, 1000) # each 1s

msgpackrpc/server.py

@@ -12,14 +12,15 @@ class Server(session.Session):
     Server is usaful for MessagePack RPC Server.
     """
 
-    def __init__(self, dispatcher, loop=None, builder=tcp):
+    def __init__(self, dispatcher, loop=None, builder=tcp, pack_encoding='utf-8', unpack_encoding=None):
         self._loop = loop or Loop()
         self._builder = builder
+        self._encodings = (pack_encoding, unpack_encoding)
         self._listeners = []
         self._dispatcher = dispatcher
 
     def listen(self, address):
-        listener = self._builder.ServerTransport(address)
+        listener = self._builder.ServerTransport(address, self._encodings)
         listener.listen(self)
         self._listeners.append(listener)
 
@@ -41,7 +42,7 @@ def on_notify(self, method, param):
 
     def dispatch(self, method, param, responder):
         try:
-            if inPy3k:
+            if inPy3k and not isinstance(method, str):
                 method = method.decode("utf-8")
             if not hasattr(self._dispatcher, method):
                 raise error.NoMethodError("'{0}' method not found".format(method))

msgpackrpc/session.py

@@ -21,7 +21,7 @@ class Session(object):
     result to the corresponding future.
     """
 
-    def __init__(self, address, timeout, loop=None, builder=tcp):
+    def __init__(self, address, timeout, loop=None, builder=tcp, reconnect_limit=5, pack_encoding='utf-8', unpack_encoding=None):
         """\
         :param address: address of the server.
         :param loop:    context object.
@@ -31,7 +31,7 @@ def __init__(self, address, timeout, loop=None, builder=tcp):
         self._loop = loop or Loop()
         self._address = address
         self._timeout = timeout
-        self._transport = builder.build_transport(self)
+        self._transport = builder.ClientTransport(self, self._address, reconnect_limit, encodings=(pack_encoding, unpack_encoding))
         self._generator = _NoSyncIDGenerator()
         self._request_table = {}
 
@@ -72,12 +72,6 @@ def on_connect_failed(self, reason):
         Called by the transport layer.
         """
 
-        def _iteritems(dic): # ugly!!!!!!
-            if inPy3k:
-                return dic.items()
-            else:
-                return dic.iteritems()
-
         # set error for all requests
         #for msgid, future in self._request_table.iteritems():
         for msgid, future in _iteritems(self._request_table):
@@ -111,7 +105,7 @@ def on_timeout(self, msgid):
 
     def step_timeout(self):
         timeouts = []
-        for msgid, future in self._request_table.iteritems():
+        for msgid, future in _iteritems(self._request_table):
             if future.step_timeout():
                 timeouts.append(msgid)
 
@@ -124,6 +118,11 @@ def step_timeout(self):
             future.set_error("Request timed out")
         self._loop.start()
 
+def _iteritems(dic): # ugly!!!!!!
+    if inPy3k:
+        return dic.items()
+    else:
+        return dic.iteritems()
 
 def _NoSyncIDGenerator():
     """\

msgpackrpc/transport/tcp.py

@@ -8,15 +8,11 @@
 from msgpackrpc import error
 
 
-def build_transport(session, reconnect_limit=5):
-    return ClientTransport(session, session.address, reconnect_limit)
-
-
 class BaseSocket(object):
-    def __init__(self, stream):
+    def __init__(self, stream, encodings):
         self._stream = stream
-        self._packer = msgpack.Packer()
-        self._unpacker = msgpack.Unpacker()
+        self._packer = msgpack.Packer(encoding=encodings[0])
+        self._unpacker = msgpack.Unpacker(encoding=encodings[1])
 
     def close(self):
         self._stream.close()
@@ -55,8 +51,8 @@ def on_notify(self, method, param):
 
 
 class ClientSocket(BaseSocket):
-    def __init__(self, stream, transport):
-        BaseSocket.__init__(self, stream)
+    def __init__(self, stream, transport, encodings):
+        BaseSocket.__init__(self, stream, encodings)
         self._transport = transport
         self._stream.set_close_callback(self.on_close)
 
@@ -78,9 +74,10 @@ def on_response(self, msgid, error, result):
 
 
 class ClientTransport(object):
-    def __init__(self, session, address, reconnect_limit):
+    def __init__(self, session, address, reconnect_limit, encodings=('utf-8', None)):
         self._session = session
         self._address = address
+        self._encodings = encodings
         self._reconnect_limit = reconnect_limit;
 
         self._connecting = 0
@@ -99,7 +96,7 @@ def send_message(self, message, callback=None):
 
     def connect(self):
         stream = IOStream(self._address.socket(), io_loop=self._session._loop._ioloop)
-        socket = ClientSocket(stream, self)
+        socket = ClientSocket(stream, self, self._encodings)
         socket.connect();
 
     def close(self):
@@ -134,8 +131,8 @@ def on_close(self, sock):
 
 
 class ServerSocket(BaseSocket):
-    def __init__(self, stream, transport):
-        BaseSocket.__init__(self, stream)
+    def __init__(self, stream, transport, encodings):
+        BaseSocket.__init__(self, stream, encodings)
         self._transport = transport
         self._stream.read_until_close(self.on_read, self.on_read)
 
@@ -150,21 +147,23 @@ def on_notify(self, method, param):
 
 
 class MessagePackServer(netutil.TCPServer):
-    def __init__(self, transport, io_loop=None):
+    def __init__(self, transport, io_loop=None, encodings=None):
         self._transport = transport
+        self._encodings = encodings
         netutil.TCPServer.__init__(self, io_loop=io_loop)
 
     def handle_stream(self, stream, address):
-        ServerSocket(stream, self._transport)
+        ServerSocket(stream, self._transport, self._encodings)
 
 
 class ServerTransport(object):
-    def __init__(self, address):
+    def __init__(self, address, encodings=('utf-8', None)):
         self._address = address;
+        self._encodings = encodings
 
     def listen(self, server):
         self._server = server;
-        self._mp_server = MessagePackServer(self, io_loop=self._server._loop._ioloop)
+        self._mp_server = MessagePackServer(self, io_loop=self._server._loop._ioloop, encodings=self._encodings)
         self._mp_server.listen(self._address.port)
 
     def close(self):

test/test_msgpackrpc.py

@@ -29,7 +29,7 @@ def _start_server(server):
         self._thread = threading.Thread(target=_start_server, args=(self._server,))
         self._thread.start()
 
-        self._client = msgpackrpc.Client(self._address)
+        self._client = msgpackrpc.Client(self._address, unpack_encoding='utf-8')
         return self._client;
 
     def tearDown(self):
@@ -40,8 +40,6 @@ def tearDown(self):
     def test_call(self):
         client = self.setup_env();
         result = client.call('hello')
-        if inPy3k:
-            result = result.decode("utf-8")
         self.assertEqual(result, "world", "'hello' result is incorrect")
 
         result = client.call('sum', 1, 2)
@@ -55,11 +53,7 @@ def test_call_async(self):
         feture1.join()
         feture2.join()
 
-        if inPy3k:
-            result = feture1.result.decode("utf-8")
-        else:
-            result = feture1.result
-        self.assertEqual(result, "world", "'hello' result is incorrect in call_async")
+        self.assertEqual(feture1.result, "world", "'hello' result is incorrect in call_async")
         self.assertEqual(feture2.result, 3, "'sum' result is incorrect in call_async")
 
     def test_notify(self):
@@ -70,16 +64,17 @@ def test_notify(self):
             client.notify('sum', 1, 2)
         except:
             result = False
-        self.assert_(result)
+        self.assertTrue(result)
 
     def test_unknown_method(self):
         client = self.setup_env();
         self.assertRaises(error.RPCError, lambda: client.call('unknown', True))
         try:
             client.call('unknown', True)
-            self.assert_(False)
+            self.assertTrue(False)
         except error.RPCError as e:
-            self.assertEqual(e.message, "'unknown' method not found", "Error message mismatched")
+            message = e.args[0]
+            self.assertEqual(message, "'unknown' method not found", "Error message mismatched")
 
 
 if __name__ == '__main__':