Source

zedify / zedify / app.py

Full commit
#!/usr/bin/env python2

from collections import defaultdict
from functools import wraps
from os.path import abspath, dirname, join
from threading import Thread, Event
import sys
import time

from gi.repository import Notify

from wx.lib.pubsub import pub
import wx
import zmq

from zedify.config import config
from zedify.version import __name__

ctx = zmq.Context()

RES_DIR = join(abspath(dirname(__file__)), 'resources')
ID_ICON_TIMER = wx.NewId()


def posix_only(f):
    @wraps(f)
    def wrapped(self, msg):
        if 'win' not in sys.environ:
            return f(self, msg)
        else:
            return None

    return wrapped


class ZedifyIcon(wx.TaskBarIcon):

    def __init__(self):
        super(ZedifyIcon, self).__init__()

        self.total_unseen = 0
        self.blink_state = False

        self.normal_icon = get_icon('tray.png')
        self.unseen_icon = get_icon('tray-unseen.png')
        self.menu = self.CreatePopupMenu()

        pub.subscribe(self.on_message_received, 'msg')

        self.set_tooltip()

        self.icon_flash = wx.Timer(self, ID_ICON_TIMER)
        wx.EVT_TIMER(self, ID_ICON_TIMER, self.blink)
        self.icon_flash.Start(750)

        self.snatch = MessageSnatcher()
        self.snatch.start()

    def CreatePopupMenu(self, counts=None):
        menu = wx.Menu()

        if counts:
            for info in counts.items():
                create_menu_item(menu, '%s:\t%s' % info, lambda ev: None)
            menu.AppendSeparator()

        create_menu_item(menu, 'Exit', self.on_exit)

        return menu

    def set_tooltip(self, counts=None):
        counts = counts or {}
        self.total_unseen = sum(counts.values())

        if self.total_unseen == 0:
            self.tip = 'No unseen messages'
            self.set_normal_icon()
        else:
            self.tip = '\n'.join(['%s: %s' % (b, c)
                                  for b, c in counts.items() if c > 0])
            self.set_unseen_icon()

    def set_normal_icon(self):
        self.set_icon(self.normal_icon)

    def set_unseen_icon(self):
        self.set_icon(self.unseen_icon)

    def set_icon(self, icon):
        self.SetIcon(icon, self.tip)

    def blink(self, evt):
        if self.total_unseen > 0:
            if self.blink_state:
                self.set_unseen_icon()
            else:
                self.set_normal_icon()

            self.blink_state = not self.blink_state
        else:
            self.blink_state = False

    def on_message_received(self, counts):
        self.set_tooltip(counts)
        self.CreatePopupMenu(counts)

    def on_exit(self, event):
        self.snatch.stop()
        wx.CallAfter(self.Destroy)


class MessageSnatcher(Thread):
    """Thread that receives and handles messages from zmq"""

    def __init__(self):
        super(MessageSnatcher, self).__init__()

        self.unseen = defaultdict(int)
        self.last_user = None
        self.last_time = None

        try:
            self.methods = config.zedify.notify.split(',')
        except AttributeError:
            self.methods = ['libnotify', 'bell']

        self.event = Event()
        self.libnotify_ready = Notify.init(__name__)

        # connect to zmq
        self.sub = ctx.socket(zmq.SUB)
        self.sub.connect(config.zedify.sub)
        self.sub.setsockopt(zmq.SUBSCRIBE, '')

    def run(self):
        """Accept messages from zmq and act upon them"""

        try:
            poller = zmq.Poller()
            poller.register(self.sub, zmq.POLLIN)

            # loop until we're asked to close the app
            while not self.event.is_set():
                socks = dict(poller.poll(1000))
                if self.sub in socks:
                    self.notify(self.sub.recv_json())

        finally:
            if self.libnotify_ready:
                Notify.uninit()

            self.sub.close()

    def notify(self, msg):
        ev = msg['event']
        handler = getattr(self, 'handle_%s' % (ev,), None)

        if handler is not None:
            now = time.time()
            if self.last_time and now - self.last_time > 15:
                self.last_user = None

            handler(msg)

            self.last_time = now

        pub.sendMessage('msg', counts=dict(self.unseen.items()))

    def handle_pm(self, msg):
        for method in self.methods:
            handler = getattr(self, '_handle_%s' % (method,), None)
            if handler:
                try:
                    handler(msg)
                except:
                    pass

        self.unseen[msg['chat']] += 1
        self.last_user = msg.get('user', None)

    @posix_only
    def _handle_libnotify(self, msg):
        if self.libnotify_ready:
            Notify.Notification.new(msg['user'], msg['msg'],
                                    'dialog-information').show()

    @posix_only
    def _handle_espeak(self, msg):
        from espeak import espeak

        text = msg['msg']
        if self.last_user != msg['user']:
            status = '%s said, "%s"' % (msg['user'], text)
        else:
            status = text

        espeak.synth(status)

    @posix_only
    def _handle_bell(self, msg):
        from subprocess import Popen
        Popen(['/usr/bin/aplay', '-q', join(RES_DIR, 'receive.wav')])

    @posix_only
    def _handle_tmux(self, msg):
        from subprocess import call

        status = '%s sent you a message' % (msg['user'],)
        call(['/usr/bin/tmux', 'display-message', status])

    def handle_mention(self, msg):
        return self.handle_pm(msg)

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

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

    def handle_buffer_switch(self, msg):
        self.unseen[msg['to_buffer']] = 0

    handle_window_switch = handle_buffer_switch

    def stop(self):
        self.event.set()


def main():
    app = wx.PySimpleApp()
    ZedifyIcon()
    app.MainLoop()


def create_menu_item(menu, label, func):
    item = wx.MenuItem(menu, -1, label)
    menu.Bind(wx.EVT_MENU, func, id=item.GetId())
    menu.AppendItem(item)

    return item


def get_icon(path):
    return wx.IconFromBitmap(wx.Bitmap(join(RES_DIR, path)))


if __name__ == '__main__':
    main()