Repository: opensistemas-hub/osbrain
Pull request #48 (2017-02-16 16:21:46)

Json serialization ( #45 )
I have been working on adding `json` serialization (related to #45  ) as well as on making the process for adding new serialization options easier.

However, there are still a few things missing:
- Docummentation and examples.
- The separator for identifying the topic and the message in PUBSUB communication is hardcoded to "b'\x80'". Perhaps it would make sense to be configurable (per agent/socket)? (Note that the separator should only used in the PUBSUB pattern when the serialization is NOT raw and there is a topic.


Patch:
diff --git a/osbrain/address.py b/osbrain/address.py
index b59da36..3c8472c 100644
--- a/osbrain/address.py
+++ b/osbrain/address.py
@@ -186,7 +186,7 @@ class AgentAddressSerializer(str):
         Serializer type (i.e.: 'raw', 'pickle'...).
     """
     def __new__(cls, value):
-        if value not in ('raw', 'pickle'):
+        if value not in ('raw', 'pickle', 'json'):
             raise ValueError('Invalid serializer type %s!' % value)
         return super().__new__(cls, value)
 
diff --git a/osbrain/agent.py b/osbrain/agent.py
index 9521b6d..bc5f316 100644
--- a/osbrain/agent.py
+++ b/osbrain/agent.py
@@ -7,6 +7,7 @@ import inspect
 import multiprocessing
 import os
 import pickle
+import json
 import signal
 import sys
 import types
@@ -28,6 +29,20 @@ from .proxy import Proxy
 from .proxy import NSProxy
 
 
+def str2bytes(message):
+    return message.encode('ascii')
+
+
+def bytes2str(message):
+    return message.decode('ascii')
+
+
+def check_type(x, check_type):
+    if not isinstance(x, check_type):
+        msg = 'Expected type `{}`, but got `{}`'.format(check_type, type(x))
+        raise TypeError(msg)
+
+
 def serialize_message(message, serializer):
     """
     Check if a message needs to be serialized and do it if that is the
@@ -46,11 +61,21 @@ def serialize_message(message, serializer):
         The serialized message, or the same message in case no
         serialization is needed.
     """
+    result = None
+
     if serializer == 'pickle':
-        return pickle.dumps(message, -1)
+        result = pickle.dumps(message, -1)
+    if serializer == 'json':
+        result = str2bytes(json.dumps(message))
     if serializer == 'raw':
-        return message
-    raise ValueError('Serializer not supported for serialization')
+        result = message
+
+    if result is None:
+        raise ValueError('Serializer not supported for serialization')
+
+    check_type(result, bytes)
+
+    return result
 
 
 def deserialize_message(message, serializer):
@@ -71,11 +96,49 @@ def deserialize_message(message, serializer):
         The deserialized message, or the same message in case no
         deserialization is needed.
     """
+    check_type(message, bytes)
+
+    result = None
+
     if serializer == 'pickle':
-        return pickle.loads(message)
+        result = pickle.loads(message)
+    if serializer == 'json':
+        result = json.loads(bytes2str(message))
     if serializer == 'raw':
-        return message
-    raise ValueError('Serializer not supported for deserialization')
+        result = message
+
+    if result is None:
+        raise ValueError('Serializer not supported for deserialization')
+
+    return result
+
+
+def compose_message(serializer, message, topic=''):
+    """
+    Compose a message and leave it ready to be sent through a socket.
+
+    Parameters
+    ----------
+    serializer : AgentAddressSerializer
+        Serialization for the message part.
+    message : anything
+        Message to be serialized. The user is the one responsible for passing
+        serializable data.
+    topic : str
+        Topic.
+
+    Returns
+    -------
+    bytes
+        The bytes representation of the final message to be sent.
+    """
+    assert isinstance(topic, str), 'Topic must be of `str` type!'
+    serialized = serialize_message(message=message, serializer=serializer)
+    topic = str2bytes(topic)
+    if topic and serializer != 'raw':
+        separator = b'\x80'
+        return topic + separator + serialized
+    return topic + serialized
 
 
 class Agent():
@@ -117,11 +180,12 @@ class Agent():
     running : bool
         Set to `True` if the agent is running (executing the main loop).
     """
-    def __init__(self, name=None, host=None):
+    def __init__(self, name=None, host=None, serializer=None):
         self.name = name
         self.host = host
         if not self.host:
             self.host = '127.0.0.1'
+        self.serializer = serializer
         self.socket = {}
         self.address = {}
         self.handler = {}
@@ -298,7 +362,7 @@ class Agent():
 
     def ping(self):
         """
-        A simple ping method testing purposes.
+        A simple ping method (for testing purposes).
         """
         return 'pong'
 
@@ -505,7 +569,9 @@ class Agent():
             'This socket requires a handler!'
         socket = self.context.socket(kind)
         transport = transport or os.environ.get('OSBRAIN_DEFAULT_TRANSPORT')
-        serializer = serializer or os.getenv('OSBRAIN_DEFAULT_SERIALIZER')
+        serializer = serializer \
+            or self.serializer \
+            or os.getenv('OSBRAIN_DEFAULT_SERIALIZER')
         addr = self._bind_socket(socket, addr=addr, transport=transport)
         server_address = AgentAddress(transport, addr, kind, 'server',
                                       serializer)
@@ -606,7 +672,7 @@ class Agent():
             handlers = {'': handlers}
         for topic in handlers.keys():
             assert isinstance(topic, str), 'Topic must be of type `str`!'
-            topic = self.str2bytes(topic)
+            topic = str2bytes(topic)
             self.socket[alias].setsockopt(zmq.SUBSCRIBE, topic)
         # Reset handlers
         self._set_handler(self.socket[alias], handlers)
@@ -754,8 +820,7 @@ class Agent():
     def _process_rep_event(self, socket_kind, socket, handler_return,
                            serializer):
         if socket_kind == 'REP' and handler_return is not None:
-            if serializer == 'pickle':
-                handler_return = pickle.dumps(handler_return, -1)
+            handler_return = serialize_message(handler_return, serializer)
             socket.send(handler_return)
 
     def _process_nonsub_event(self, socket_kind, socket, serialized):
@@ -786,12 +851,32 @@ class Agent():
                                 serializer)
 
     def _process_sub_message(self, serializer, message):
+        """
+        Return the received message in a PUBSUB communication.
+
+        Parameters
+        ----------
+        message : bytes
+            Received message without any treatment. Note that we do not know
+            whether there is a topic or not.
+
+        Returns
+        -------
+        anything
+            The content of the message passed.
+        """
+        separator = b'\x80'
+
         if serializer == 'pickle':
-            # Since SUB messages might have a `topic`, we need to be careful to
-            # only deserialize the non-topic part. Pickle objects always start
-            # with b'\x80', so the pure message part will start from that char.
-            sepp = message.index(b'\x80')
-            message = memoryview(message)[sepp:]
+            if not message.startswith(separator):
+                sepp = message.index(separator) + 1
+                message = bytes(memoryview(message)[sepp:])
+
+        if serializer == 'json':
+            if separator in message:
+                sepp = message.index(separator) + 1
+                message = bytes(memoryview(message)[sepp:])
+
         return deserialize_message(message=message, serializer=serializer)
 
     def _process_sub_event(self, socket, serialized):
@@ -812,7 +897,7 @@ class Agent():
         message = self._process_sub_message(serializer, serialized)
 
         for str_topic in handlers:
-            btopic = self.str2bytes(str_topic)
+            btopic = str2bytes(str_topic)
             if not serialized.startswith(btopic):
                 continue
             # Call the handler (with or without the topic)
@@ -848,30 +933,41 @@ class Agent():
 
         return serializer
 
-    def str2bytes(self, message):
-        return message.encode('ascii')
-
     def send(self, address, message, topic=''):
         """
-        TODO
+        Send a message through the specified address.
+
+        Note that replies in a REQREP pattern do not use this function in
+        order to be sent.
         """
-        assert isinstance(topic, str), 'Topic must be of `str` type!'
         serializer = self.address[address].serializer
-        serialized = serialize_message(message=message, serializer=serializer)
-        topic = self.str2bytes(topic)
-        self.socket[address].send(topic + serialized)
+        message = compose_message(serializer, message, topic)
+        self.socket[address].send(message)
 
     def recv(self, address):
         """
-        TODO
+        Receive a message from the specified address.
 
         This method is only used in REQREP communication patterns.
+
+        Parameters
+        ----------
+        address :
+
+        Returns
+        -------
+        anything
+            The content received in the address.
+
         """
         message = self.socket[address].recv()
         serializer = self.address[address].serializer
         return deserialize_message(message=message, serializer=serializer)
 
     def send_recv(self, address, message):
+        """
+        This method is only used in REQREP communication patterns.
+        """
         self.send(address, message)
         return self.recv(address)
 
@@ -927,7 +1023,8 @@ class AgentProcess(multiprocessing.Process):
     Agent class. Instances of an Agent are system processes which
     can be run independently.
     """
-    def __init__(self, name, nsaddr=None, addr=None, base=Agent):
+    def __init__(self, name, nsaddr=None, addr=None, serializer=None,
+                 base=Agent):
         super().__init__()
         self.name = name
         self.daemon = None
@@ -935,6 +1032,7 @@ class AgentProcess(multiprocessing.Process):
         if self.port is None:
             self.port = 0
         self.nsaddr = nsaddr
+        self.serializer = serializer
         self.base = base
         self.shutdown_event = multiprocessing.Event()
         self.queue = multiprocessing.Queue()
@@ -952,7 +1050,8 @@ class AgentProcess(multiprocessing.Process):
             return
         self.queue.put('STARTED')
 
-        self.agent = self.base(name=self.name, host=self.host)
+        self.agent = self.base(name=self.name, host=self.host,
+                               serializer=self.serializer)
         uri = self.daemon.register(self.agent)
         ns.register(self.name, uri)
         ns.release()
@@ -999,7 +1098,7 @@ class AgentProcess(multiprocessing.Process):
         self.kill()
 
 
-def run_agent(name, nsaddr=None, addr=None, base=Agent):
+def run_agent(name, nsaddr=None, addr=None, base=Agent, serializer=None):
     """
     Ease the agent creation process.
 
@@ -1022,7 +1121,8 @@ def run_agent(name, nsaddr=None, addr=None, base=Agent):
     """
     if not nsaddr:
         nsaddr = os.environ.get('OSBRAIN_NAMESERVER_ADDRESS')
-    AgentProcess(name, nsaddr=nsaddr, addr=addr, base=base).start()
+    AgentProcess(name, nsaddr=nsaddr, addr=addr, base=base,
+                 serializer=serializer).start()
     proxy = Proxy(name, nsaddr)
     proxy.run()
     return proxy


Test patch:
diff --git a/osbrain/tests/test_agent.py b/osbrain/tests/test_agent.py
index 62cac1f..72d0f8b 100644
--- a/osbrain/tests/test_agent.py
+++ b/osbrain/tests/test_agent.py
@@ -173,6 +173,44 @@ def test_socket_creation(nsaddr):
     assert 'alias2' in addresses
 
 
+def test_correct_serialization(nsaddr):
+    """
+    Test that the right serializer is being used when using the different
+    initialization options.
+    """
+    # Without specifying a serializer
+    a0 = run_agent('a0')
+    addr0 = a0.bind('PUB', 'alias0')
+    _address = a0.get_attr('address')
+    assert _address[addr0].serializer \
+        == os.getenv('OSBRAIN_DEFAULT_SERIALIZER')
+
+    # Specifying a serializer at Agent level
+    # Test is twice with different serializer to avoid collision with the
+    # environment variable
+    a1 = run_agent('a1', serializer='raw')
+    addr1 = a1.bind('PUB', 'alias1')
+    _address = a1.get_attr('address')
+    assert _address[addr1].serializer == 'raw'
+
+    a2 = run_agent('a2', serializer='pickle')
+    addr2 = a2.bind('PUB', 'alias2')
+    _address = a2.get_attr('address')
+    assert _address[addr2].serializer == 'pickle'
+
+    # Specifying a serializer at Socket level
+    # Test is twice with different serializer to avoid collision with the
+    # environment variable
+    a3 = run_agent('a3')
+    addr3 = a3.bind('PUB', 'alias3', serializer='raw')
+    _address = a3.get_attr('address')
+    assert _address[addr3].serializer == 'raw'
+    a4 = run_agent('a4')
+    addr4 = a4.bind('PUB', 'alias4', serializer='json')
+    _address = a4.get_attr('address')
+    assert _address[addr4].serializer == 'json'
+
+
 def test_reqrep(nsaddr):
     """
     Simple request-reply pattern between two agents.
diff --git a/osbrain/tests/test_agent_pubsub_topics.py b/osbrain/tests/test_agent_pubsub_topics.py
index 6c8b29a..50b35bd 100644
--- a/osbrain/tests/test_agent_pubsub_topics.py
+++ b/osbrain/tests/test_agent_pubsub_topics.py
@@ -82,6 +82,78 @@ def test_pubsub_topics_pickle(nsaddr):
     assert message_04 in a5.get_attr('received_list')
 
 
+def test_pubsub_topics_json(nsaddr):
+    """
+    Simple publisher-subscriber pattern test.
+
+    Different messages sent with different agents subscribed to different
+    topics are tested within this method.
+    """
+    a0 = run_agent('a0')
+    a1 = run_agent('a1')
+    a2 = run_agent('a2')
+    a3 = run_agent('a3')
+    a4 = run_agent('a4')
+    a5 = run_agent('a5')
+
+    for agent in (a1, a2, a3, a4, a5):
+        agent.set_attr(received_list=[])
+
+    addr = a0.bind('PUB', alias='pub', serializer='json')
+
+    a1.connect(addr, handler=log_received_to_list)
+    a2.connect(addr, handler={'foo': log_received_to_list})
+    a3.connect(addr, handler={'bar': log_received_to_list,
+                              'foo': log_received_to_list})
+    a4.connect(addr, handler={'bar': log_received_to_list})
+    a5.connect(addr, handler={'fo': log_received_to_list})
+
+    # Give some time for all the agents to connect
+    time.sleep(0.1)
+
+    # Send some messages
+    message_01 = 'Hello'
+    a0.send('pub', message_01)
+
+    message_02 = 'World'
+    a0.send('pub', message_02, topic='foo')
+
+    message_03 = 'FOO'
+    a0.send('pub', message_03, topic='foobar')
+
+    message_04 = 'BAR'
+    a0.send('pub', message_04, topic='fo')
+
+    # Give some time for all the agents to handle the message
+    time.sleep(0.1)
+
+    # Check each agent received the corresponding messages
+    assert message_01 in a1.get_attr('received_list')
+    assert message_02 in a1.get_attr('received_list')
+    assert message_03 in a1.get_attr('received_list')
+    assert message_04 in a1.get_attr('received_list')
+
+    assert message_01 not in a2.get_attr('received_list')
+    assert message_02 in a2.get_attr('received_list')
+    assert message_03 in a2.get_attr('received_list')
+    assert message_04 not in a2.get_attr('received_list')
+
+    assert message_01 not in a3.get_attr('received_list')
+    assert message_02 in a3.get_attr('received_list')
+    assert message_03 in a3.get_attr('received_list')
+    assert message_04 not in a3.get_attr('received_list')
+
+    assert message_01 not in a4.get_attr('received_list')
+    assert message_02 not in a4.get_attr('received_list')
+    assert message_03 not in a4.get_attr('received_list')
+    assert message_04 not in a4.get_attr('received_list')
+
+    assert message_01 not in a5.get_attr('received_list')
+    assert message_02 in a5.get_attr('received_list')
+    assert message_03 in a5.get_attr('received_list')
+    assert message_04 in a5.get_attr('received_list')
+
+
 def test_pubsub_topics_raw(nsaddr):
     """
     Simple publisher-subscriber pattern test.
diff --git a/osbrain/tests/test_agent_serialization.py b/osbrain/tests/test_agent_serialization.py
index 54ad92e..b838a8a 100644
--- a/osbrain/tests/test_agent_serialization.py
+++ b/osbrain/tests/test_agent_serialization.py
@@ -1,5 +1,6 @@
 import time
 import pickle
+import json
 
 import zmq
 import pytest
@@ -7,6 +8,7 @@ import pytest
 from osbrain import run_agent
 from osbrain.agent import serialize_message
 from osbrain.agent import deserialize_message
+from osbrain.agent import compose_message
 
 from common import nsaddr  # pragma: no flakes
 from common import nsproxy  # pragma: no flakes
@@ -16,15 +18,62 @@ def set_received(agent, message, topic=None):
     agent.received = message
 
 
+def test_message_composer():
+    msg_1 = b'Chain of bytes'
+    msg_2 = [1, 3, "Hello"]
+    topic = "test topic"
+    topic_bytes = b"test topic"
+
+    separator = b'\x80'
+
+    # Test raw serialization composing
+    assert compose_message('raw', msg_1) == b'Chain of bytes'
+    assert compose_message('raw', msg_1, topic) == b'test topicChain of bytes'
+
+    # Test pickle serialization composing
+    assert compose_message('pickle', msg_1) == pickle.dumps(msg_1, -1)
+    assert compose_message('pickle', msg_2) == pickle.dumps(msg_2, -1)
+
+    assert compose_message('pickle', msg_1, topic) \
+        == topic_bytes + separator + pickle.dumps(msg_1, -1)
+    assert compose_message('pickle', msg_2, topic) \
+        == topic_bytes + separator + pickle.dumps(msg_2, -1)
+
+    # Test json serialization composing
+    msg_1_str = 'Chain of bytes'
+
+    assert compose_message('json', msg_1_str) \
+        == json.dumps(msg_1_str).encode('ascii')
+    assert compose_message('json', msg_2) \
+        == json.dumps(msg_2).encode('ascii')
+
+    assert compose_message('json', msg_1_str, topic) \
+        == topic_bytes + separator + json.dumps(msg_1_str).encode('ascii')
+    assert compose_message('json', msg_2, topic) \
+        == topic_bytes + separator + json.dumps(msg_2).encode('ascii')
+
+
 def test_serialize_message():
     """
     Test basic serialization.
     """
+    # Raw serialization
     test = b'asdf'
     assert test == serialize_message(message=test, serializer='raw')
+
+    # Pickle serialization
     test = [0, 1]
     assert test == pickle.loads(serialize_message(message=test,
                                 serializer='pickle'))
+
+    # Json serialization
+    assert test == json.loads(serialize_message(message=test,
+                              serializer='json').decode('ascii'))
+    # Un-serializable type
+    with pytest.raises(TypeError):
+        serialize_message(message=b'Hello', serializer='json')
+
+    # Incorrect serializer
     with pytest.raises(ValueError):
         serialize_message(message=test, serializer='foo')
 
@@ -33,11 +82,29 @@ def test_deserialize_message():
     """
     Test basic deserialization.
     """
+    # Raw deserialization
     test = b'asdf'
     assert test == deserialize_message(message=test, serializer='raw')
+
+    # Pickle deserialization
     test = [0, 1]
     assert test == deserialize_message(message=pickle.dumps(test, -1),
                                        serializer='pickle')
+
+    # Json deserialization
+    assert test == \
+        deserialize_message(message=json.dumps(test).encode('ascii'),
+                            serializer='json')
+
+    # Wrong type
+    with pytest.raises(TypeError):
+        deserialize_message(message='Hello', serializer='raw')
+    with pytest.raises(TypeError):
+        deserialize_message(message=[1, 2, 3], serializer='pickle')
+    with pytest.raises(TypeError):
+        deserialize_message(message=None, serializer='json')
+
+    # Incorrect serializer
     with pytest.raises(ValueError):
         deserialize_message(message=b'x', serializer='foo')
 
@@ -72,6 +139,21 @@ def test_reqrep_pickle(nsaddr):
     assert response == 'OK'
 
 
+def test_reqrep_json(nsaddr):
+    """
+    Simple request-reply pattern between two agents with json serialization.
+    """
+    def rep_handler(agent, message):
+        return 'OK'
+
+    a0 = run_agent('a0')
+    a1 = run_agent('a1')
+    addr = a0.bind('REP', 'reply', rep_handler, serializer='json')
+    a1.connect(addr, 'request')
+    response = a1.send_recv('request', 'Hello world')
+    assert response == 'OK'
+
+
 def test_reqrep_raw_zmq_outside(nsaddr):
     """
     Simple request-reply pattern between an agent and a direct ZMQ connection.
@@ -131,6 +213,22 @@ def test_pushpull_pickle(nsaddr):
     assert a1.get_attr('received') == message
 
 
+def test_pushpull_json(nsaddr):
+    """
+    Simple push-pull pattern test with json serialization.
+    """
+    a0 = run_agent('a0')
+    a1 = run_agent('a1')
+    a1.set_attr(received=None)
+    addr = a1.bind('PULL', handler=set_received, serializer='json')
+    a0.connect(addr, 'push')
+    message = 'Hello world'
+    a0.send('push', message)
+    while not a1.get_attr('received'):
+        time.sleep(0.01)
+    assert a1.get_attr('received') == message
+
+
 def test_pushpull_raw_zmq_outside(nsaddr):
     """
     Simple push-pull pattern test. Channel without serialization.
@@ -191,6 +289,22 @@ def test_pubsub_pickle(nsaddr):
     assert a1.get_attr('received') == message
 
 
+def test_pubsub_json(nsaddr):
+    """
+    Simple publisher-subscriber pattern test with json serialization.
+    """
+    a0 = run_agent('a0')
+    a1 = run_agent('a1')
+    a1.set_attr(received=None)
+    addr = a0.bind('PUB', alias='pub', serializer='json')
+    a1.connect(addr, handler=set_received)
+    message = 'Hello world'
+    while not a1.get_attr('received'):
+        a0.send('pub', message)
+        time.sleep(0.1)
+    assert a1.get_attr('received') == message
+
+
 def test_pubsub_raw_zmq_outside(nsaddr):
     """
     Simple publisher-subscriber pattern test. Channel without serialization.