Simplify package structure

tfw/internals/networking/__init__.py

@@ -0,0 +1,4 @@
+from .serialization import serialize_tfw_msg, deserialize_tfw_msg, with_deserialize_tfw_msg, message_bytes
+from .server_connector import ServerUplinkConnector, ServerDownlinkConnector, ServerConnector
+from .event_handler_connector import EventHandlerConnector
+from .scope import Scope

tfw/internals/networking/event_handler_connector.py

@@ -0,0 +1,48 @@
+import logging
+
+import zmq
+from zmq.eventloop.zmqstream import ZMQStream
+
+from .serialization import serialize_tfw_msg, with_deserialize_tfw_msg
+
+LOG = logging.getLogger(__name__)
+
+
+class EventHandlerDownlinkConnector:
+    def __init__(self, bind_addr):
+        self._zmq_pull_socket = zmq.Context.instance().socket(zmq.PULL)
+        self._zmq_pull_socket.setsockopt(zmq.RCVHWM, 0)
+        self._zmq_pull_stream = ZMQStream(self._zmq_pull_socket)
+        self._zmq_pull_socket.bind(bind_addr)
+        LOG.debug('Pull socket bound to %s', bind_addr)
+
+    def register_callback(self, callback):
+        callback = with_deserialize_tfw_msg(callback)
+        self._zmq_pull_stream.on_recv(callback)
+
+    def close(self):
+        self._zmq_pull_stream.close()
+
+
+class EventHandlerUplinkConnector:
+    def __init__(self, bind_addr):
+        self._zmq_pub_socket = zmq.Context.instance().socket(zmq.PUB)
+        self._zmq_pub_socket.setsockopt(zmq.SNDHWM, 0)
+        self._zmq_pub_socket.bind(bind_addr)
+        LOG.debug('Pub socket bound to %s', bind_addr)
+
+    def send_message(self, message: dict):
+        self._zmq_pub_socket.send_multipart(serialize_tfw_msg(message))
+
+    def close(self):
+        self._zmq_pub_socket.close()
+
+
+class EventHandlerConnector(EventHandlerDownlinkConnector, EventHandlerUplinkConnector):
+    def __init__(self, downlink_bind_addr, uplink_bind_addr):
+        EventHandlerDownlinkConnector.__init__(self, downlink_bind_addr)
+        EventHandlerUplinkConnector.__init__(self, uplink_bind_addr)
+
+    def close(self):
+        EventHandlerDownlinkConnector.close(self)
+        EventHandlerUplinkConnector.close(self)

tfw/internals/networking/scope.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class Scope(Enum):
+    ZMQ = 'zmq'
+    WEBSOCKET = 'websocket'
+    BROADCAST = 'broadcast'

tfw/internals/networking/serialization.py

@@ -0,0 +1,104 @@
+"""
+TFW JSON message format
+
+message:
+{
+    "key":     string,    # addressing
+    "data":    {...},     # payload
+    "trigger": string     # FSM trigger
+}
+
+ZeroMQ's sub-pub sockets use enveloped messages
+(http://zguide.zeromq.org/page:all#Pub-Sub-Message-Envelopes)
+and TFW also uses them internally. This means that on ZMQ sockets
+we always send the messages key separately and then the actual
+message (which contains the key as well) like so:
+
+socket.send_multipart([message['key'], message])
+
+The purpose of this module is abstracting away this low level behaviour.
+"""
+
+import json
+from functools import wraps
+
+
+def serialize_tfw_msg(message):
+    """
+    Create TFW multipart data from message dict
+    """
+    return _serialize_all(message['key'], message)
+
+
+def with_deserialize_tfw_msg(fun):
+    @wraps(fun)
+    def wrapper(message_parts):
+        message = deserialize_tfw_msg(*message_parts)
+        return fun(message)
+    return wrapper
+
+
+def deserialize_tfw_msg(*args):
+    """
+    Return message from TFW multipart data
+    """
+    return _deserialize_all(*args)[1]
+
+
+def _serialize_all(*args):
+    return tuple(
+        _serialize_single(arg)
+        for arg in args
+    )
+
+
+def _deserialize_all(*args):
+    return tuple(
+        _deserialize_single(arg)
+        for arg in args
+    )
+
+
+def _serialize_single(data):
+    """
+    Return input as bytes
+    (serialize input if it is JSON)
+    """
+    if not isinstance(data, str):
+        data = message_bytes(data)
+    return _encode_if_needed(data)
+
+
+def message_bytes(message):
+    return json.dumps(message, sort_keys=True).encode()
+
+
+def _deserialize_single(data):
+    """
+    Try parsing input as JSON, return it as
+    string if parsing fails.
+    """
+    try:
+        return json.loads(data)
+    except ValueError:
+        return _decode_if_needed(data)
+
+
+def _encode_if_needed(value):
+    """
+    Return input as bytes
+    (encode if input is string)
+    """
+    if isinstance(value, str):
+        value = value.encode('utf-8')
+    return value
+
+
+def _decode_if_needed(value):
+    """
+    Return input as string
+    (decode if input is bytes)
+    """
+    if isinstance(value, (bytes, bytearray)):
+        value = value.decode('utf-8')
+    return value

tfw/internals/networking/server_connector.py

@@ -0,0 +1,65 @@
+import logging
+
+import zmq
+from zmq.eventloop.zmqstream import ZMQStream
+
+from .scope import Scope
+from .serialization import serialize_tfw_msg, with_deserialize_tfw_msg
+
+LOG = logging.getLogger(__name__)
+
+
+class ServerDownlinkConnector:
+    def __init__(self, connect_addr):
+        self.keys = []
+        self._on_recv_callback = None
+        self._zmq_sub_socket = zmq.Context.instance().socket(zmq.SUB)
+        self._zmq_sub_socket.setsockopt(zmq.RCVHWM, 0)
+        self._zmq_sub_socket.connect(connect_addr)
+        self._zmq_sub_stream = ZMQStream(self._zmq_sub_socket)
+
+    def subscribe(self, *keys):
+        for key in keys:
+            self._zmq_sub_socket.setsockopt_string(zmq.SUBSCRIBE, key)
+            self.keys.append(key)
+
+    def unsubscribe(self, *keys):
+        for key in keys:
+            self._zmq_sub_socket.setsockopt_string(zmq.UNSUBSCRIBE, key)
+            self.keys.remove(key)
+
+    def register_callback(self, callback):
+        self._on_recv_callback = callback
+        self._zmq_sub_stream.on_recv(with_deserialize_tfw_msg(self._on_recv))
+
+    def _on_recv(self, message):
+        key = message['key']
+        if key in self.keys or '' in self.keys:
+            self._on_recv_callback(message)
+
+    def close(self):
+        self._zmq_sub_stream.close()
+
+
+class ServerUplinkConnector:
+    def __init__(self, connect_addr):
+        self._zmq_push_socket = zmq.Context.instance().socket(zmq.PUSH)
+        self._zmq_push_socket.setsockopt(zmq.SNDHWM, 0)
+        self._zmq_push_socket.connect(connect_addr)
+
+    def send_message(self, message, scope=Scope.ZMQ):
+        message['scope'] = scope.value
+        self._zmq_push_socket.send_multipart(serialize_tfw_msg(message))
+
+    def close(self):
+        self._zmq_push_socket.close()
+
+
+class ServerConnector(ServerDownlinkConnector, ServerUplinkConnector):
+    def __init__(self, downlink_connect_addr, uplink_connect_addr):
+        ServerDownlinkConnector.__init__(self, downlink_connect_addr)
+        ServerUplinkConnector.__init__(self, uplink_connect_addr)
+
+    def close(self):
+        ServerDownlinkConnector.close(self)
+        ServerUplinkConnector.close(self)