pypsyc/mjacob2/pypsyc/core/psyc.py

"""
    pypsyc.core.psyc
    ~~~~~~~~~~~~~~~~

    Classes which are related to the entity layer.

    :copyright: 2010 by Manuel Jacob
    :license: MIT
"""
import logging
from collections import defaultdict
from os import urandom

from greenlet import greenlet

from pypsyc.core.mmp import Header
from pypsyc.util import Waiter


log = logging.getLogger(__name__)


def _get_binary_arg(it, must_len, first_fragment):
    fragments = [first_fragment]
    cur_len = len(first_fragment)
    while cur_len < must_len:
        f = it.next()
        fragments.append(f)
        cur_len += len(f) + 1
    return '\n'.join(fragments)

class PSYCPacket(object):
    """
    Initialize a new PSYCPacket with the given arguments.

    :param modifiers: entity modifiers
    :type modifiers: dict
    :param mc: method
    :type mc: string or None
    :param data: data
    :type data: string

    Attributes:

    .. attribute:: modifiers

       entity modifiers as described `here <http://about.psyc.eu/Spec:Packet>`_

    .. attribute:: cvars

       same as :attr:`modifiers`\[':'\]

    .. attribute:: mc

       method as described `here <http://about.psyc.eu/Spec:Packet>`_

    .. attribute:: data

       data as described `here <http://about.psyc.eu/Spec:Packet>`_
    """

    def __init__(self, modifiers=None, mc=None, data=''):
        self.modifiers = modifiers or {}
        self.cvars = self.modifiers.setdefault(':', {})
        self.mc = mc
        self.data = data

    @classmethod
    def from_kwds(cls, mc=None, data='', **kwds):
        """
        Return a new PSYCPacket which is initialized with mc and data given by
        arguments and cvars given by keyword arguments.

        ::

            packet1 = PSYCPacket.from_kwds(mc='_message_public',
                                           data="hello", _somevar='foo')

            packet2 = PSYCPacket(mc='_message_public', data="hello",
                                 modifiers={':': {'_somevar': 'foo'}})

            packet1 == packet2 # True

        :param mc: method
        :type mc: string or None
        :param data: data
        :type data: string
        :rtype: PSYCPacket
        """
        return cls({':': kwds}, mc, data)

    def render(self):
        """Return a generator of lines which represent this packet."""
        for glyph, vars in self.modifiers.iteritems():
            for x, y in vars.iteritems():
                if '\n' in y:
                    yield '%s%s %s\t%s' % (glyph, x, len(y), y)
                else:
                    yield '%s%s\t%s' % (glyph, x, y)
        if self.mc:
            yield self.mc
            if self.data:
                yield self.data
        else:
            assert not self.data

    @classmethod
    def parse(cls, content):
        """Parse a list of lines to a packet."""
        modifiers = defaultdict(dict)
        it = iter(content)
        for line in it:
            if line.startswith('_'): # body
                return cls(modifiers, line, '\n'.join(it))
            else: # entity-modifier
                glyph = line[0]
                x, y = line[1:].split('\t', 1)
                if ' ' in x: # binary-arg
                    var, l = x.split(' ')
                    modifiers[glyph][var] = _get_binary_arg(it, int(l), y)
                else: # simple-arg
                    modifiers[glyph][x] = y
        return cls(modifiers, None, '')

    def __repr__(self):
        return '\n'.join((
            "PSYCPacket(",
            "    modifiers={",
            ",\n".join(' '*8 + "'%s': %s" % i for i in self.modifiers.items()),
            "    },",
            "    mc=%r,",
            "    data=%r",
            ")"
        )) % (self.mc, self.data)

    def __eq__(self, other):
        return (self.modifiers == other.modifiers and
                self.mc == other.mc and
                self.data == other.data)

    def __ne__(self, other):
        return not self.__eq__(other)


class PSYCObject(object):
    """
    This is the base class for all classes that provide PSYC functionality,
    such as receiving and sending :class:`PSYCPacket`\s.

    Subclassed by :class:`pypsyc.server.Entity`. Used by
    :class:`pypsyc.server.routing.ServerCircuit`.
    """

    def __init__(self, sendfunc, uni=None):
        self._send = sendfunc
        self.uni = uni
        self.handlers = {}
        self.tags = {}

    def handle_packet(self, header, content):
        packet = PSYCPacket.parse(content)
        packet.header = header

        tag_relay = header.get('_tag_relay')
        if tag_relay:
            self.tags.pop(tag_relay).callback(packet)
            return packet

        mc = packet.mc
        while mc:
            f = self.handlers.get(mc)
            if f is not None:
                break
            mc = mc.rpartition('_')[0]
        else:
            return packet

        tag = header.get('_tag')
        if tag is None:
            try:
                f(packet)
            except Exception:
                log.exception("error calling handler %s of %s", mc, self)
        else:
            def respond():
                try:
                    response = f(packet)
                except Exception:
                    log.exception("error calling handler %s of %s", mc, self)
                    response = PSYCPacket(mc='_error_internal')
                self.sendmsg(header.source, response, {'_tag_relay': tag})
            greenlet(respond).switch()
        return packet

    def sendmsg(self, target, packet=None, header=None, **kwds):
        """
        Send a packet.

        1.
        Simple example usage::

            packet = PSYCPacket(mc='_message_public',
                                modifiers={':': {'_somevar': 'foo'}})
            psyc.sendmsg(target, packet)

        This simply sends the packet to the target.


        2.
        The following code block is equal to the previous::

            psyc.sendmsg(target, mc='_message_public', _somevar='foo')

        See also :meth:`PSYCPacket.from_kwds`.

        :param target: target of the packet
        :type target: :class:`pypsyc.core.mmp.Uni` or \
                      :class:`pypsyc.core.mmp.Circuit`
        :param packet: packet to send, use keyword arguments if None
        :type packet: :class:`PSYCPacket` or None
        :param header: additional header variables
        :type header: dict or None
        :rtype: PSYCPacket or None
        """
        if packet is None:
            packet = PSYCPacket.from_kwds(**kwds)

        header = Header() if header is None else Header(header)
        if self.uni is not None:
            header['_source'] = self.uni

        if packet.mc and packet.mc.startswith('_request'):
            tag = urandom(20).replace('\n', '')
            header['_tag'] = tag
            waiter = self.tags[tag] = Waiter()
        else:
            waiter = None

        if isinstance(target, str):
            header['_target'] = target
            header._init()
            self._send(header, packet.render())
        else: # the target is a circuit
            target.send(header, packet.render())

        if waiter is not None:
            return waiter.get()

    def add_handler(self, handler):
        self.handlers.update((i[6:], getattr(handler, i)) for i in
                             dir(handler) if i.startswith('handle_'))

    def __repr__(self):
        return '<PSYCObject uni=%s>' % self.uni