swayswitch/src/swayswitch

#!/usr/bin/env python3
"""A simple windowswitcher for sway"""
# switch_window, get_windows and extract_nodes_iterative is derived from
# https://github.com/tobiaspc/wofi-scripts Copyright (c) 2020 Tobi which
# is covered by the MIT License. However this program is licensed under
# the GPLv3+ Copyright (c) 2020 Micke Nordin <hej@mic.ke>.
# See LICENSE file for details.

import json
import math
import os
import subprocess
from io import BytesIO

import cairosvg
import toml
import wx


class SwaySwitch(wx.Frame):  # pylint: disable=no-member
    """Frame for the swayswitcher"""
    def __init__(self, *args, **kw):  # pylint: disable=unused-argument
        """Constructor"""
        wx.Frame.__init__(self, None, title="", style=wx.STAY_ON_TOP)  # pylint: disable=no-member
        # create a panel in the frame
        self.pnl = wx.Panel(self)  # pylint: disable=no-member

        # Some xdg data
        self.home = os.environ.get('HOME')
        self.base_dirs = self.get_base_dirs()
        # get windows from sway
        windows = get_windows()

        # Config defaults
        self.label_len = 20
        self.icon_size = 128
        # Can be overwritten in config_file $HOME/.local/swayswitch/config
        self.set_config_from_file()

        self.icon_dirs = self.get_icon_dirs()

        # and create a sizer to manage the layout of child widgets
        x_and_y = int(math.sqrt(len(windows)) + 0.5)
        self.sizer = wx.GridSizer(x_and_y)  # pylint: disable=no-member
        self.pnl.SetSizer(self.sizer)
        self.set_buttons(windows)
        self.sizer.Fit(self)
        self.sizer.Layout()

    def get_base_dirs(self):
        """We want to follow the XDG standard if possible"""
        base_dirs = []
        xdg_dirs = os.environ.get('XDG_DATA_DIRS')
        pos_dirs = []
        if xdg_dirs:
            pos_dirs = xdg_dirs.split(":")
        if not pos_dirs:
            pos_dirs = [
                "/usr/share", "/usr/local/share", self.home + "/.local/share",
                self.home + "/.local/share/flatpak/exports/share",
                "/var/lib/flatpak/exports/share"
            ]
        for pos_dir in pos_dirs:
            if os.path.exists(pos_dir):
                base_dirs.append(pos_dir)
        return base_dirs

    def get_desktop_file(self, command):
        """From here we return first dektopfile"""
        desktop_file = None
        for base_dir in self.base_dirs:
            directory = base_dir + "/applications"
            if os.path.exists(directory):
                command = "grep -s -l {} {}/*.desktop".format(
                    command, directory)
                for data in [i for i in run_command(command) if i]:
                    for result in data.decode().split('\n'):
                        desktop_file = result
                        break
        return desktop_file

    def get_icon_bitmap(self, icon):
        """Create a bitmap with right size from svg or png"""
        if icon:
            if icon.endswith(".svg"):
                svgpng = cairosvg.svg2png(
                    bytestring=open(icon).read().encode('utf-8'))
                image = wx.Image(BytesIO(svgpng), wx.BITMAP_TYPE_PNG)  # pylint: disable=no-member
            else:
                unscaled_bitmap = wx.Bitmap(icon)  # pylint: disable=no-member
                image = unscaled_bitmap.ConvertToImage()
            image = image.Scale(self.icon_size, self.icon_size,
                                wx.IMAGE_QUALITY_HIGH)  # pylint: disable=no-member
            bitmap = wx.Bitmap(image)  # pylint: disable=no-member
        else:
            bitmap = wx.Bitmap()  # pylint: disable=no-member
        return bitmap

    def get_icon_path(self, desktop_file):
        """Glue function to get icon path"""
        icon_name = get_icon_name(desktop_file)
        pos_icons = self.get_pos_icons(icon_name)
        icon_path = self.get_right_size_icon(pos_icons)
        return icon_path

    def get_icon_dirs(self):
        """Find icon dirs"""
        icon_dirs = []
        for base_dir in [
                directory for directory in self.base_dirs if directory
        ]:
            directory = base_dir + "/icons"
            icon_dirs.append(directory)
        icon_dirs.append(self.home + "/.icons")
        icon_dirs.append("/usr/share/pixmaps")

        return icon_dirs

    def get_label(self, window):
        """Construct the label for the window"""
        try:
            label = window['window_properties']['class']
        except KeyError:
            try:
                label = window['app_id']
                if len(label) < self.label_len:
                    label = label + ':' + window['name'][:self.label_len -
                                                         len(label)]
            except KeyError:
                label = window['name']
            if len(label) > self.label_len:
                label = label[:self.label_len]

        label = "ws" + str(window['workspace']) + ": " + label
        return label

    def get_pos_icons(self, icon_name):
        """Collect possible icons, prefer svg"""
        possible_icons = []
        for icon_dir in [
                directory for directory in self.icon_dirs if directory
        ]:
            if os.path.exists(icon_dir):
                command = "find {} -name *{}.svg".format(icon_dir, icon_name)
                result = run_command(command)[0].split()
                if not result:
                    command = "find {} -name *{}.png".format(
                        icon_dir, icon_name)
                    result = run_command(command)[0].split()
                for data in result:
                    icon_cand = data.decode()
                    if icon_cand not in possible_icons:
                        possible_icons.append(icon_cand)
        return possible_icons

    def get_right_size_icon(self, possible_icons):
        """Try to find prefered size"""
        icon = None
        for pos_icon in possible_icons:
            if pos_icon.endswith(".svg"):
                icon = pos_icon
                break
            if not icon and (str(self.icon_size) + "x" + str(self.icon_size)
                             in pos_icon):
                icon = pos_icon
        if not icon:
            try:
                icon = sorted(possible_icons)[0]
            except IndexError:
                icon = ""
        return icon

    def on_key_press(self, event):
        """Intercept esc key press"""
        keycode = event.GetUnicodeKey()
        if keycode == wx.WXK_ESCAPE:  # pylint: disable=no-member
            # enter normal mode and exit
            command = 'swaymsg mode "default"'
            run_command(command)
            self.Close(True)
        else:
            event.Skip(True)

    def set_buttons(self, windows):
        """Loop over windows and create buttons for them"""
        for window in windows:
            # This is setting up an inner sizer with a static text label and an image icon
            inner_sizer = wx.BoxSizer(orient=wx.VERTICAL)  # pylint: disable=no-member
            label = self.get_label(window)
            winid = window['id']
            size = wx.Window.GetFont(self).GetPointSize() * self.label_len  # pylint: disable=no-member
            command = get_command(window['pid'])
            desktop_file = self.get_desktop_file(command)
            icon_path = self.get_icon_path(desktop_file)
            bitmap = self.get_icon_bitmap(icon_path)
            btn = wx.BitmapButton(  # pylint: disable=no-member
                self.pnl,
                id=winid,
                bitmap=bitmap,
                size=wx.Size(size, size))  # pylint: disable=no-member
            btn.Bind(
                wx.EVT_BUTTON,
                lambda event, mwinid=winid: self.switch_window(event, mwinid))
            # Set up esc keybinding
            self.Bind(wx.EVT_CHAR_HOOK, lambda event: self.on_key_press(event))  # pylint: disable=unnecessary-lambda
            statictext = wx.StaticText(self.pnl, -1, label)  # pylint: disable=no-member
            inner_sizer.Add(statictext, 0, wx.ALIGN_CENTER)  # pylint: disable=no-member
            inner_sizer.Add(btn, 0, wx.ALIGN_CENTER)  # pylint: disable=no-member
            inner_sizer.Fit(self)
            inner_sizer.Layout()
            self.sizer.Add(inner_sizer, 0, wx.ALIGN_CENTER)  # pylint: disable=no-member

    def set_config_from_file(self):
        """If user has created a config file we will use values from there"""
        config_file = self.home + "/.local/swayswitch/config"
        if os.path.exists(config_file):
            try:
                cfg = toml.load(config_file)
                if cfg["label_len"]:
                    self.label_len = cfg["label_len"]
                if cfg["icon_size"]:
                    self.icon_size = cfg["icon_size"]
            except toml.TomlDecodeError:
                #If use has formating errors we warn to stderr, but move on with defaults
                os.write(
                    2, b"WARNING: formatting errors in " +
                    config_file.encode() + b"\n")

    def switch_window(self, event, winid):  # pylint: disable=unused-argument
        """Switches the focus to the given id and enter normalmode"""
        command = 'swaymsg [con_id={}] focus, mode "default"'.format(winid)
        run_command(command)
        self.Close(True)


def get_command(pid):
    """Returns a list of all json window objects"""
    command = "ps h c -p {} -o command".format(pid)
    result = run_command(command)[0].rstrip().decode()
    return result


def get_icon_name(desktop_file):
    """ Find icon from a desktop file"""
    command = "grep Icon= {}".format(desktop_file)
    for data in [i.rstrip() for i in run_command(command) if i]:
        try:
            icon_name = data.decode().split('=')[1]
            return icon_name
        except IndexError:
            return ""


def get_windows():
    """Returns a list of all json window objects"""
    command = "swaymsg -t get_tree"
    data = json.loads(run_command(command)[0])

    # Select outputs that are active
    windows = []
    for output in data['nodes']:

        # The scratchpad (under __i3) is not supported
        if output.get('name') != '__i3' and output.get('type') == 'output':
            workspaces = output.get('nodes')
            for workspace in workspaces:
                if workspace.get('type') == 'workspace':
                    windows += extract_nodes_iterative(workspace)
    return windows


def extract_nodes_iterative(workspace):
    """Extracts all windows from a sway workspace json object"""
    all_nodes = []

    floating_nodes = workspace.get('floating_nodes')

    for floating_node in floating_nodes:
        floating_node['workspace'] = workspace['num']
        all_nodes.append(floating_node)

    nodes = workspace.get('nodes')

    for node in nodes:

        # Leaf node
        if len(node.get('nodes')) == 0:
            node['workspace'] = workspace['num']
            all_nodes.append(node)
        # Nested node, handled iterative
        else:
            for inner_node in node.get('nodes'):
                inner_node['workspace'] = workspace['num']
                nodes.append(inner_node)

    return all_nodes


def run_command(command):
    """Run a command on system and capture result"""
    process = subprocess.Popen(command,
                               shell=True,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    result = process.communicate()
    return result


# Entry point
if __name__ == "__main__":
    # When this module is run (not imported) then create the app, the
    # frame, show it, and start the event loop.
    app = wx.App()
    frm = SwaySwitch(None, title="")
    frm.Show()
    app.MainLoop()