Source

zedify / zedify / i3.py

"""
Weechat notification plugin for py3status.

Notifications received from zedifyd via 0mq will be used to update the i3
status bar to show unseen messages. Notifications go away when you switch to a
buffer with unseen messages or when you type into such a buffer.
"""

from threading import Thread, Event
import atexit
import os
import re
import sh
import site
import time

import zmq

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


SHORT_RE = re.compile(r'(\w.*?)[A-Z_]')
SND_PATH = os.path.join(site.getsitepackages()[0], 'zedify',
                        'resources', 'receive.wav')
REFRESH_INTERVAL = int(config.zedifier.get('interval', 10))
log = get_logger('zedifier')


ctx = zmq.Context()

SUB = ctx.socket(zmq.SUB)
SUB.setsockopt_unicode(zmq.SUBSCRIBE, '')
SUB.connect(config.zedifier.sub)

PUSH = ctx.socket(zmq.PUSH)
PUSH.connect(config.zedifier.push)

MESSAGES = {}


class Py3status:

    def __init__(self):

        worker = None
        event = Event()

        def handle(msg):
            global MESSAGES

            log.info('Received message: %s' % (msg,))
            ev = msg.get('event', None)

            if ev is None:
                return

            if ev == 'current_status':
                MESSAGES = msg.get('unseen', {})
            elif ev in ('pm', 'mention'):
                b = msg.get('chat', None)
                log.info('Got PM/mention event: %s' % (b,))
                if b is None:
                    return

                x = sh.notify_send(b, msg['msg'], _bg=True)
                x = sh.pacmd('play-file', SND_PATH, 0, _bg=True)

                assert x

        def receive(ev):
            log.info('About to receive messages')

            last_update = time.time()
            p = zmq.Poller()
            p.register(SUB, zmq.POLLIN)

            request_status(True)

            try:
                while not ev.is_set():
                    # request a status update if we haven't received anything
                    # recently
                    if time.time() - last_update >= REFRESH_INTERVAL:
                        request_status()

                    try:
                        socks = dict(p.poll(1000))
                        if SUB in socks:
                            last_update = time.time()
                            handle(SUB.recv_json())
                    except:
                        log.exception('Error')

            finally:
                SUB.close()

        def request_status(full=False):
            log.info('Requesting unseen count update')
            PUSH.send_json({'event': 'status', 'full': full})

        def stop_receiving(signal=None, frame=None):
            log.warn('Attempting to terminate zedifier')
            try:
                log.warn('Setting event...')
                event.set()

                log.warn('Joining...')
                worker.join()

                log.warn('Closing socket...')

            except:
                log.exception('Error cleaning up')

            log.info('Done.')

        atexit.register(stop_receiving)
        worker = Thread(target=receive, args=[event])
        worker.start()

    def unseen(self, json, cfg):
        """Show any unread weechat messages in the status bar"""

        def short(s):
            m = SHORT_RE.match(s)
            if m:
                return m.group(1)
            else:
                return s

        unseen = ['{}:{}'.format(short(u), c)
                  for u, c in MESSAGES.items()]

        res = {
            'full_text': '|'.join(unseen),
            'name': 'zedifier_unseen',
        }

        if cfg and 'color_degraded' in cfg:
            res['color'] = cfg['color_degraded']

        return (0, res)


if __name__ == '__main__':
    Py3status().unseen(None, None)