Source

crossroads / crossroads / message.py

# -*- coding: utf-8 -*-
'''Crossroads.IO messages.'''

from stuf.six import PY3, tobytes, tounicode, b
from ctypes import byref, sizeof, string_at, c_ubyte

from . import lowest as xs


class BaseMsg(object):

    def __init__(self, size, source=None):
        self.source = source
        self.more = self.rc = self.dest = None

    def __bytes__(self, *args, **kwargs):
        if self.dest is not None:
            return tobytes(string_at(self.ref, len(self)), 'latin-1')
        return b('')

    def __unicode__(self, *args, **kwargs):
        return tounicode(tobytes(self, 'utf-8'))

    __str__ = __unicode__ if PY3 else __bytes__

    @property
    def array(self):
        return bytearray(self.__bytes__())


class BaseMultipart(object):

    def __init__(self, size, source=None):
        self.source = source
        self.more = self.rc = self.dest = None
        self.parts = []
        self.__add__ = self.parts.append

    def __bytes__(self, *args, **kwargs):
        if self.dest is not None:
            return b('').join(tobytes(
                string_at(byref(i), sizeof(i)), 'latin-1') for i in self.parts)
        return b('')

    def __unicode__(self, *args, **kwargs):
        return tounicode(tobytes(self, 'utf-8'))

    __str__ = __unicode__ if PY3 else __bytes__

    @property
    def dest(self):
        return self.parts[-1]

    @property
    def array(self):
        return bytearray(self.__bytes__())


class Message(BaseMsg):

    __slots__ = 'source dest last_rc more'.split()

    def __init__(self, size, source=None):
        super(Message, self).__init__(size, source)
        if source is None:
            self.dest = xs.array(c_ubyte, size)

    def __len__(self):
        return len(self.source) if self.dest is None else sizeof(self.dest)

    @property
    def ref(self):
        return None if self.dest is None else byref(self.dest)


class Multipart(BaseMultipart):

    def __init__(self, size, source=None):
        super(Multipart, self).__init__(size, source)

    def __len__(self):
        return len(self.source) if self.dest is None else sizeof(self.dest)

    @property
    def ref(self):
        return None if self.dest is None else byref(self.dest)


class XSMessage(BaseMsg):

    __slots__ = 'source dest last_rc more ref'.split()

    def __init__(self, data=None):
        super(XSMessage, self).__init__(32, data)
        self.dest = self.ref = dest = xs.msg_t()
        if data is None:
            self.last_rc = xs.msg_init(dest)
        else:
            self.last_rc = xs.msg_init_size(dest, 32)
            xs.msg_init_size(dest, 32)
            xs.memmove(xs.msg_data(dest), data, 32)

    def __len__(self):
        return sizeof(self.dest)

    def close(self):
        self.last_rc = xs.msg_close(self.dest)


class XSMultipart(BaseMultipart):

    def __init__(self, size, source=None):
        super(XSMultipart, self).__init__(size, source)

    def __len__(self):
        return sizeof(self.dest)

    def close(self):
        self.last_rc = sum(xs.msg_close(i) for i in self.parts)