messages.py 1.2 KB
Newer Older
1 2
import re

3 4
import msgpack

Robin Obůrka's avatar
Robin Obůrka committed
5
from .exceptions import InvalidMsgError, InvalidMsgTypeError
6

7
SN_MSG_REGEXP = "^([a-z0-9_]+/)*[a-z0-9_]+$"
8 9
SN_MSG = re.compile(SN_MSG_REGEXP)

Robin Obůrka's avatar
Robin Obůrka committed
10

11 12 13 14 15 16
def parse_msg(data):
    """ Gets a Sentinel-type ZMQ message and parses message type and its
    payload.
    """
    try:
        msg_type = str(data[0], encoding="UTF-8")
17 18
        if not SN_MSG.match(msg_type):
            raise InvalidMsgTypeError("Bad message type definition")
19 20 21 22 23
        payload = msgpack.unpackb(data[1], encoding="UTF-8")

    except IndexError:
        raise InvalidMsgError("Not enough parts in message")

24 25 26
    except (TypeError, msgpack.exceptions.UnpackException, UnicodeDecodeError):
        raise InvalidMsgError("Broken message")

27 28 29 30 31 32 33
    return msg_type, payload


def encode_msg(msg_type, data):
    """ Gets string message type and its's string data. Then, both of them are
    packed to be prepared for zmg.send_multipart().
    """
34 35 36 37 38 39 40

    if not SN_MSG.match(msg_type or ""):
        raise InvalidMsgTypeError("Bad message type definition")

    if not data or type(data) != dict:
        raise InvalidMsgError("Empty payload parameter")

41 42 43 44
    b = bytes(msg_type, encoding="UTF-8")
    msg = msgpack.packb(data)

    return (b, msg)