Source

zedify / zedify / remote.py

Full commit
from collections import defaultdict

import zmq

from zedify.logger import get_logger
from zedify.config import config


class ZedifyDaemon:

    def __init__(self):
        self.unseen = defaultdict(list)
        self.log = get_logger('zedifyd')

        ctx = zmq.Context()

        self.pub = ctx.socket(zmq.PUB)
        self.pub.bind(config.zedifyd.pub)

        self.sub = ctx.socket(zmq.PULL)
        self.sub.bind(config.zedifyd.pull)

    def serve(self):
        """
        Simple distribution function.

        Receives zmq messages and blasts them out to any interested parties.
        """

        try:
            while True:
                msg = self.sub.recv_json()
                self.pub.send_json(msg)

                self.log.info('Received message: {}'.format(msg))
                ev = msg.get('event', None)
                handler = getattr(self, 'handle_{}'.format(ev), None)
                if handler is not None:
                    handler(msg)
                else:
                    self.log.error('Unexpected message: {}'.format(msg))

                if ev != 'status':
                    self.handle_status()

        except KeyboardInterrupt:
            pass
        finally:
            self.pub.close()
            self.sub.close()

    def handle_status(self, msg):
        """Publish the number of unseen messages"""

        self.pub.send_json({
            'event': 'current_status',
            'unseen': {k: len(v) for k, v in self.unseen.items() if v}
        })

        if msg.get('full', False):
            [self.pub.send_json(m)
                for msgs in self.unseen.values()
                for m in msgs]

    def handle_typing(self, msg):
        """Reset unseen count for the buffer we're typing into"""

        self.unseen[msg['in_buffer']] = []

    def handle_buffer_switch(self, msg):
        """Reset unseen count for the buffer we're switching into"""

        self.unseen[msg['to_buffer']] = []

    def handle_mention(self, msg):
        """Increment the unseen count for the buffer"""

        self.handle_pm(msg)

    def handle_pm(self, msg):
        """Increment the unseen count for the buffer"""

        self.unseen[msg['chat']].append(msg)


def main():
    ZedifyDaemon().serve()


if __name__ == '__main__':
    main()