main.py

import os
import logging
import gi

gi.require_version('Gtk', '3.0')
gi.require_version('Notify', '0.7')

from locale import atof, setlocale, LC_NUMERIC
from gi.repository import Notify
from itertools import islice
from subprocess import check_output, check_call, CalledProcessError
from ulauncher.api.client.Extension import Extension
from ulauncher.api.client.EventListener import EventListener
from ulauncher.api.shared.event import KeywordQueryEvent, ItemEnterEvent
from ulauncher.api.shared.item.ExtensionSmallResultItem import ExtensionSmallResultItem
from ulauncher.api.shared.action.RenderResultListAction import RenderResultListAction
from ulauncher.api.shared.action.ExtensionCustomAction import ExtensionCustomAction


logger = logging.getLogger(__name__)
kill_icon = 'images/kill.png'
pause_icon = 'images/pause.png'
play_icon = 'images/play.png'
exec_icon = 'images/executable.png'
dead_icon = 'images/dead.png'


class ProcessKillerExtension(Extension):

    def __init__(self):
        super(ProcessKillerExtension, self).__init__()
        self.subscribe(KeywordQueryEvent, KeywordQueryEventListener())
        self.subscribe(ItemEnterEvent, ItemEnterEventListener())
        setlocale(LC_NUMERIC, '')  # set to OS default locale;

    def show_notification(self, title, text=None, icon=kill_icon):
        logger.debug('Show notification: %s' % text)
        icon_full_path = os.path.join(os.path.dirname(__file__), icon)
        Notify.init("KillerExtension")
        Notify.Notification.new(title, text, icon_full_path).show()


class KeywordQueryEventListener(EventListener):

    def on_event(self, event, extension):
        return RenderResultListAction(list(islice(self.generate_results(event), 15)))

    def generate_results(self, event):
        for (pid, cpu, cmd) in get_process_list():
            name = '[%s%% CPU] %s' % (cpu, cmd) if cpu > 1 else cmd
            on_enter = {'alt_enter': False, 'pid': pid, 'cmd': cmd}
            on_alt_enter = on_enter.copy()
            on_alt_enter['alt_enter'] = True
            if event.get_argument():
                if event.get_argument() in cmd:
                    yield ExtensionSmallResultItem(icon=exec_icon,
                                                   name=name,
                                                   on_enter=ExtensionCustomAction(on_enter),
                                                   on_alt_enter=ExtensionCustomAction(on_alt_enter, keep_app_open=True))
            else:
                yield ExtensionSmallResultItem(icon=exec_icon,
                                               name=name,
                                               on_enter=ExtensionCustomAction(on_enter),
                                               on_alt_enter=ExtensionCustomAction(on_alt_enter, keep_app_open=True))


class ItemEnterEventListener(EventListener):

    def kill(self, extension, pid, signal):
        cmd = ['kill', '-s', signal, pid]
        logger.info(' '.join(cmd))

        try:
            check_call(cmd) == 0
            icon = ""
            if signal == "STOP":
                text = "It's paused now"
            elif signal == "CONT":
                text = "It's alive now"
            else:
                text = "It's dead now"
                icon = dead_icon
            extension.show_notification("Done", text, icon=icon)
        except CalledProcessError as e:
            extension.show_notification("Error", "'kill' returned code %s" % e.returncode)
        except Exception as e:
            logger.error('%s: %s' % (type(e).__name__, e))
            extension.show_notification("Error", "Check the logs")
            raise

    def show_signal_options(self, data):
        result_items = []
        options = [
            ('TERM', '15 TERM (default)', kill_icon),
            ('KILL', '9 KILL', kill_icon),
            ('HUP', '1 HUP', kill_icon),
            ('STOP', '19 STOP (pause)', pause_icon),
            ('CONT', '18 CONT (resume)', play_icon),
        ]
        for sig, name, icon in options:
            on_enter = data.copy()
            on_enter['alt_enter'] = False
            on_enter['signal'] = sig
            result_items.append(ExtensionSmallResultItem(icon=icon,
                                                         name=name,
                                                         highlightable=False,
                                                         on_enter=ExtensionCustomAction(on_enter)))
        return RenderResultListAction(result_items)

    def on_event(self, event, extension):
        data = event.get_data()
        if data['alt_enter']:
            return self.show_signal_options(data)
        else:
            self.kill(extension, data['pid'], data.get('signal', 'TERM'))


def get_process_list():
    """
    Returns a list of tuples (PID, %CPU, COMMAND)
    """
    env = os.environ.copy()
    env['COLUMNS'] = '200'
    out = check_output(['ps', '-eo', 'pid,%cpu,cmd', '--sort', '-%cpu'], env=env).decode('utf8')
    for line in out.split('\n'):
        col = line.split()
        try:
            int(col[0])
        except (ValueError, IndexError):
            # not a number
            continue

        pid = col[0]
        cpu = atof(col[1])
        cmd = ' '.join(col[2:])
        if 'top -bn' in cmd:
            continue

        yield (pid, cpu, cmd)


if __name__ == '__main__':
    ProcessKillerExtension().run()