window cli - manage windows from the command line

window.py

#!/usr/bin/env python3

import logging
import re
import shlex
import time
from argparse import ArgumentParser
from functools import reduce
from math import sqrt
from os import path, devnull as DEVNULL
from subprocess import Popen

import gi
import yaml

gi.require_version('Gdk', '3.0')
gi.require_version('GdkX11', '3.0')
gi.require_version('Gtk', '3.0')
gi.require_version('Wnck', '3.0')
from gi.repository import Gdk, GdkX11, Gtk, Wnck

LOG = logging.getLogger(path.basename(__file__))
FORMATTER = logging.Formatter('%(asctime)s %(name)s:%(funcName)s:%(lineno)d %(levelname)s: %(message)s')
STREAM_HANDLER = logging.StreamHandler()
STREAM_HANDLER.setLevel(logging.INFO)
STREAM_HANDLER.setFormatter(FORMATTER)
LOG.setLevel(logging.INFO)
LOG.addHandler(STREAM_HANDLER)


class WindowManager(object):
    def __init__(self, config):
        """
        constructor
        :param config: yaml config file
        """
        self.CONFIG = yaml.load(config)
        # options
        self.CLOSE_OPT = 'close_windows'
        self.GROUP_OPT = 'group'
        self.CENTER_OPT = 'center'
        self.MAX_OPT = 'maximize'
        self.DISTR_OPT = 'distribute'
        # terminals are treated differently in certain situations
        self.TTY = 'terminal'
        # list format
        self.FORMAT = self.CONFIG['format']
        # window geometry
        self.GEOM = self.CONFIG['geometry']
        # wnck configuration
        self.WNCK_FLAGS = Wnck.WindowMoveResizeMask.X | \
                          Wnck.WindowMoveResizeMask.Y | \
                          Wnck.WindowMoveResizeMask.WIDTH | \
                          Wnck.WindowMoveResizeMask.HEIGHT
        self.WNCK_GRAVITY = Wnck.WindowGravity.NORTHWEST
        # wnck init
        self.screen = None
        self.windows = None
        self.init_wnck()
        # apps that increment configured workspace numbers
        # they are always on workspace 1 and configured by setting workspace number to 'offset'
        self.offset_apps = []
        self.offset = 0
        # app configurations
        self.apps = self.CONFIG['apps']
        self.app_names = []
        self.app_aliases = {}
        self.win_title_patterns = {}
        self.init_app_config()
        # launch control
        self.launch_count = 0
        self.launching_apps = []
        self.launched_apps = []

    def init_wnck(self):
        self.screen = Wnck.Screen.get_default()
        # wait...
        while Gtk.events_pending():
            Gtk.main_iteration()

        self.windows = self.screen.get_windows()

    def init_app_config(self):
        """
        initializes app configurations
        """
        for app_name in self.apps:
            app_config = self.apps[app_name]
            if 'alias' in app_config:
                self.app_aliases[app_config['alias']] = app_name
            if 'pattern' in app_config:
                self.win_title_patterns[app_name] = app_config['pattern']
            if app_config['workspace'] < 0:
                self.offset_apps.append(app_name)
                self.apps[app_name]['workspace'] = 0  # evaluates to 1

        self.app_names = list(self.apps.keys())

    def get_app_name(self, win):
        """
        returns a canonized form of the freedesktop name which is sometimes a bit strange
        :param win: Wnck window
        :return: canonized string
        """
        app_name = None
        candidate1 = win.get_application().get_name()
        LOG.debug('app name: "{}"'.format(candidate1))
        candidate1 = re.sub(r'(\w+\.){2,}(\w+)', r'\2', candidate1.lower().strip())

        candidate2 = win.get_name()
        LOG.debug('window title: "{}"'.format(candidate2))
        candidate2 = candidate2.lower().strip()

        if candidate1 in self.app_names:
            app_name = candidate1
        elif candidate2 in self.app_names:
            app_name = candidate2
        else:
            for name, pattern in self.win_title_patterns.items():
                if re.search(pattern, candidate2):
                    app_name = name

        return self.translate_aliases(app_name if app_name is not None else candidate1)

    def translate_aliases(self, app_name):
        """
        translates aliases into app-names
        :param app_name: app_name or alias
        :return str
        """
        return self.app_aliases[app_name] if app_name in self.app_aliases else app_name

    def is_app_launched(self, app_name):
        """
        checks if an instance of an app is already launched.
        :param app_name: name of the app.
        :return boolean
        """
        return len(self.get_windows_by_app_name(app_name)) > 0

    def get_windows_by_app_name(self, app_name):
        """
        returns a list of Wnck win objects of the given application.
        :param app_name: str
        :return [Wnck window]
        """
        app_name = self.translate_aliases(app_name)
        return [x for x in self.windows if self.get_app_name(x) == app_name]

    def get_windows_by_regex(self, regex, inactive_wins_only=False):
        """
        returns a list of Wnck win objects with app_names and/or titles matching the given regex.
        :param regex: regular expression
        :param inactive_wins_only: restrict to inactive windows
        :return [Wnck window]
        """
        matcher = re.compile(self.translate_aliases(regex), re.I)
        wins = []
        for win in self.get_all_windows(inactive_wins_only=inactive_wins_only):
            if matcher.search('{} {}'.format(self.get_app_name(win), win.get_name())):
                wins.append(win)
        return wins

    def get_all_windows(self, active_workspace_only=False, inactive_wins_only=False):
        """
        returns a list of all windows only filtered by state and workspace
        :param active_workspace_only: restrict to windows of the active workspace
        :param inactive_wins_only: restrict to inactive windows
        :return: [Wnck window]
        """
        wins = []
        if active_workspace_only or inactive_wins_only:
            for win in self.windows:
                if (not win.is_active() or not inactive_wins_only) \
                        and (not active_workspace_only or win.is_on_workspace(self.screen.get_active_workspace())):
                    wins.append(win)
        else:
            wins = self.windows
        return wins

    def get_windows(self, option, args):
        """
        selects windows according to given options and additional commandline arguments
        :param option: desired option or function
        :param args: commandline arguments
        :return [Wnck window]
        """
        windows = []
        inactive_wins_only = option in [self.CLOSE_OPT, self.MAX_OPT]

        # all windows
        if args['all']:
            windows = self.get_all_windows(inactive_wins_only=inactive_wins_only)
        # all windows on the active workspace
        elif args['all_workspace']:
            windows = self.get_all_windows(active_workspace_only=True, inactive_wins_only=inactive_wins_only)
        # all matching windows.
        elif args[option]:
            for arg in args[option]:
                windows.extend(self.get_windows_by_regex(arg, inactive_wins_only=inactive_wins_only))
        # active window
        else:
            windows.append(self.screen.get_active_window())

        return windows

    def sort_windows(self, windows, app_name=True, workspace=True, terminal=True, active=True, reverse=False):
        """
        sorts the given list of windows (stable sort)
        :param windows: [Wnck window]
        :param app_name: sort by app_name
        :param workspace: sort by workspace
        :param terminal: put terminals on top
        :param active: put active on top
        :param reverse: reverse order
        :return: sorted [Wnck window]
        """
        wins = windows[:]

        if app_name:
            wins.sort(key=lambda x: self.get_app_name(x), reverse=reverse)
        if workspace:
            wins.sort(key=lambda x: x.get_workspace().get_number(), reverse=reverse)
        if terminal:
            wins.sort(key=lambda x: self.get_app_name(x).replace(self.TTY, '0'), reverse=reverse)
        if active:
            wins.sort(key=lambda x: x.is_active(), reverse=not reverse)

        return wins

    def get_apps_to_be_opened(self, aliases, default=False):
        """
        translates aliases to application names, checks offsets and optionally appends default apps.
        :param aliases: application aliases as provided in the config file
        :param default: include default apps
        :return: [str]
        """
        app_names = []
        try:
            app_names = [self.app_aliases[x] for x in aliases]
        except KeyError as e:
            LOG.error('unknown app alias {}'.format(e))

        # also include default apps
        if default:
            app_names.extend([app_name for app_name in self.app_names if self.apps[app_name]['default']])

        # set possible offset
        LOG.debug('launched:')
        offset_candidates = [self.get_app_name(x) for x in self.get_all_windows()]
        offset_candidates.extend(app_names)
        if any([x in self.offset_apps for x in offset_candidates]):
            self.offset = 1

        return app_names

    def open_apps(self, args, default=False, set_geometry=False):
        """
        opens new windows of the given application and moves them to predefined workspaces.
        :param args: list of application shortcuts which are to be launched
        :param default: switch for including default applications.
        :param set_geometry: switch for setting default geometries
        default value is False.
        """
        self.screen.connect('window-opened', self.on_window_opened, self.screen.get_active_window(), set_geometry)

        # launch
        for app_name in self.get_apps_to_be_opened(args, default=default):
            if not self.is_app_launched(app_name):
                LOG.debug('launching {}'.format(app_name))
                self.launching_apps.append(app_name)
                self.launch_count += 1
                execstr = path.expanduser(self.apps[app_name]['exec'])
                with open(DEVNULL) as DNULL:
                    Popen(shlex.split(execstr), stdout=DNULL, stderr=DNULL)
            else:
                LOG.debug('{} is already launched'.format(app_name))

        # wait for opening windows
        if self.launch_count > 0:
            LOG.debug('gtk main start')
            Gtk.main()

    def on_window_opened(self, screen, window, active, set_geometry):
        """
        callback. listens for new windows and moves them to the predefined workspaces
        :param sreen: Wnck screen
        :param window: Wnck window which is to be moved
        :param active: Wnck window which is to be activated after moving
        :param set_geometry: switch for setting default geometries
        """
        app_name = self.get_app_name(window)
        LOG.debug('"{}" window opened'.format(app_name))
        if app_name in self.launching_apps:
            LOG.debug('"{}" will be moved and its geometry set'.format(app_name))
            offset = self.offset if app_name != self.TTY else 0
            self.move_window(
                window,
                screen.get_workspace(self.apps[app_name]['workspace'] + offset),
                active
            )
            if set_geometry:
                self.set_default_window_geometry(window, app_name)
            # don't decrement launch_count if an app opens another window
            if app_name not in self.launched_apps:
                self.launched_apps.append(app_name)
                self.launch_count -= 1
        # exit the mainloop after all windows are opened and moved.
        if self.launch_count == 0:
            LOG.debug('gtk main quit (level: {})'.format(Gtk.main_level()))
            Gtk.main_quit()
            time.sleep(1)
            active.activate(self.get_current_event_time())

    def rearrange_windows(self, set_geometry=False):
        """
        moves open windows to predefined workspaces and sets their default geometry
        :param set_geometry: switch for setting default geometries
        """
        wins = []
        # look for offset apps and collect configured apps
        for win in self.get_all_windows():
            app_name = self.get_app_name(win)
            if app_name in self.apps:
                if app_name in self.offset_apps:
                    self.offset = 1
                wins.append([app_name, win])
        # rearrange configured app windows
        for win in wins:
            app_name, win = win
            offset = self.offset if app_name != self.TTY else 0
            self.move_window(
                win,
                self.screen.get_workspace(self.apps[app_name]['workspace'] + offset),
                self.screen.get_active_window()
            )
            if set_geometry:
                self.set_default_window_geometry(win, app_name)

    def set_default_window_geometry(self, window, app_name):
        """
        sets the default geometry of the given window
        :param window: Wnck window
        :param app_name: name of the app the windows belongs to
        """
        default_geometry = self.apps[app_name]['geometry']
        if default_geometry == self.MAX_OPT:
            self.maximize_windows(windows=[window])
        elif default_geometry == self.CENTER_OPT:
            self.center_windows(windows=[window])
        elif default_geometry == self.GROUP_OPT:
            self.cascade_windows(windows=[window], move_to_active=False)
        else:
            LOG.error('unknown geometry "{}" for application "{}"'.format(default_geometry, app_name))

    def focus_window(self, name):
        """
        activate the window with a title matching the given regex
        :param name: regex for app_name or window title
        """
        wins = self.get_windows_by_regex(name)
        if wins:
            # only one window is possible
            if len(wins) > 1:
                LOG.warning('found more than one window. setting focus on first...')
            win = wins[0]
            workspace = win.get_workspace()
            workspace.activate(self.get_current_event_time())
            win.activate(self.get_current_event_time())

    def move_window(self, window, workspace_nr, active):
        """
        moves the given window to the given workspace and activates the given window.
        :param window: Wnck window which is to be moved
        :param workspace_nr: number of the target workspace
        :param active: Wnck window which is to be activated after moving
        """
        # open in background
        window.make_below()
        window.move_to_workspace(workspace_nr)
        window.unmake_below()
        active.activate(self.get_current_event_time())

    def close_apps(self, apps=None, args=None):
        """
        closes all windows of the given app
        :param apps: [str]
        :param args: commandline arguments
        """
        apps = apps if apps else args['close_apps']
        for app_name in apps:
            for win in self.get_windows_by_app_name(app_name):
                if not win.is_active():
                    LOG.debug('closing "{}"'.format(win.get_name()))
                    win.close(0)

    def close_windows(self, windows=None, args=None):
        """
        closes the defined windows.
        :param windows: [Wnck window]
        :param args: commandline arguments
        """
        windows = windows if windows else self.get_windows(self.CLOSE_OPT, args)
        for win in windows:
            LOG.debug('closing "{}"'.format(win.get_name()))
            win.close(0)

    def cascade_windows(self, windows=None, args=None, move_to_active=True):
        """
        cascades the given windows
        :param windows: [Wnck window]
        :param args: commandline arguments
        :param move_to_active: move windows to the active workspace
        """
        active = self.screen.get_active_window()
        windows = windows if windows else self.get_windows(self.GROUP_OPT, args)
        origin = int(self.GEOM['origin'])
        offset = int(self.GEOM['offset'])
        i = 0

        for win in self.sort_windows(windows, workspace=False, reverse=True):
            new_origin = origin + i * offset
            if move_to_active:
                win.move_to_workspace(self.screen.get_active_workspace())
            win.unmaximize()
            win.set_geometry(
                self.WNCK_GRAVITY,
                self.WNCK_FLAGS,
                new_origin,
                new_origin + int(self.GEOM['menubar']),
                int(self.screen.get_width() * float(self.GEOM['dimx'])),
                int(self.screen.get_height() * float(self.GEOM['dimy']))
            )
            i += 1

        active.activate(self.get_current_event_time())

    def distribute_windows(self, windows=None, args=None):
        """
        distributes the given windows across the screen
        :param windows: [Wnck window]
        :param args: commandline arguments
        """
        active = self.screen.get_active_window()
        windows = windows if windows else self.get_windows(self.DISTR_OPT, args)

        if windows:
            # calculate distribution parameters
            win_count = len(windows)
            rows = int(sqrt(win_count))
            cols = int(win_count / rows + win_count % rows)
            if (rows - 1) * cols == win_count:
                rows -= 1
            x = 0
            y = 0
            dim_x = self.screen.get_width() / cols
            dim_y = self.screen.get_height() / rows
            i = 0

            for win in self.sort_windows(windows, workspace=False, reverse=True):
                win.move_to_workspace(self.screen.get_active_workspace())
                win.unmaximize()
                win.set_geometry(self.WNCK_GRAVITY, self.WNCK_FLAGS, x, y, dim_x, dim_y)
                i += 1
                if i < cols:
                    x += dim_x
                else:
                    y += dim_y
                    x = 0
                    i = 0

        active.activate(self.get_current_event_time())

    def center_windows(self, windows=None, args=None):
        """
        centers the given windows
        :param windows: [Wnck window]
        :param args: commandline arguments
        """
        active = self.screen.get_active_window()
        windows = windows if windows else self.get_windows(self.CENTER_OPT, args)
        dim_x = int(self.screen.get_width() * float(self.GEOM['dimx']))
        dim_y = int(self.screen.get_height() * float(self.GEOM['dimy']))
        x = int((self.screen.get_width() - dim_x) / 2)
        y = int((self.screen.get_height() - dim_y - int(self.GEOM['menubar'])) / 2)

        for win in windows:
            win.unmaximize()
            win.set_geometry(self.WNCK_GRAVITY, self.WNCK_FLAGS, x, y, dim_x, dim_y)

        active.activate(self.get_current_event_time())

    def maximize_windows(self, windows=None, args=None):
        """
        maximizes the given windows
        :param windows: [Wnck window]
        :param args: commandline arguments
        """
        active = self.screen.get_active_window()
        windows = windows if windows else self.get_windows(self.MAX_OPT, args)

        for win in windows:
            win.unmaximize()
            win.set_geometry(
                self.WNCK_GRAVITY,
                self.WNCK_FLAGS,
                0,
                0,
                self.screen.get_width(),
                self.screen.get_height()
            )
            win.maximize()

        active.activate(self.get_current_event_time())

    def move_workspaces(self, workspace_numbers):
        """
        move workspaces.
        2 arg: switch two workspaces; 
        1 arg: insert empty workspace at the given position; 
        0 arg: remove gaps between workspaces
        :param workspace_numbers: workspaces to switch | position to insert | emptiness
        """
        windows = self.get_all_windows(inactive_wins_only=True)

        if len(workspace_numbers) == 0:
            # close empty workspaces
            non_empty = 0
            move = 0
            for win in self.sort_windows(windows, workspace=True):
                workspace_number = win.get_workspace().get_number()
                delta = workspace_number - non_empty
                if delta > 1 or workspace_number == move:
                    if workspace_number != move:
                        non_empty = workspace_number - delta + 1
                        move = workspace_number
                    self.move_window(win, self.screen.get_workspace(non_empty), self.screen.get_active_window())
                elif workspace_number != move:
                    non_empty = workspace_number

        elif len(workspace_numbers) == 1:
            # insert an empty workspace
            for win in windows:
                workspace_number = win.get_workspace().get_number()
                if workspace_number >= workspace_numbers[0]:
                    self.move_window(
                        win,
                        self.screen.get_workspace(workspace_number + 1),
                        self.screen.get_active_window()
                    )

        elif len(workspace_numbers) == 2:
            # switch two workspaces
            for win in windows:
                workspace_number = win.get_workspace().get_number()
                if workspace_number in workspace_numbers:
                    target = workspace_numbers[0] if workspace_number == workspace_numbers[1] else workspace_numbers[1]
                    self.move_window(
                        win,
                        self.screen.get_workspace(target),
                        self.screen.get_active_window()
                    )

    def print_window_list(self):
        """
        lists all open windows.
        """
        workspace_header = 'Workspace'
        application_header = 'Application'
        # extract and sort window information
        windows = [[x.get_workspace().get_number(), self.get_app_name(x), x.get_name()] for x in
                   self.sort_windows(self.get_all_windows(), terminal=False, active=False)]
        # define list format
        max_app_len = reduce(lambda x, y: max(x, y), [len(x[1]) for x in windows + [['', application_header, '']]])
        max_title_len = int(self.FORMAT['maxtitle'])
        spacing = int(self.FORMAT['spacing'])
        # header
        print('\n{0}{1}{2}\n'.format(
            workspace_header.ljust(len(workspace_header) + spacing),
            application_header.ljust(max_app_len + spacing),
            'Window Title'
        ))
        # windows
        for win in windows:
            title = win[2].lower()
            print('{0}{1}{2}{3}'.format(
                str(win[0]).center(len(workspace_header)),
                ' '.ljust(spacing),
                win[1].ljust(max_app_len + spacing),
                title if len(title) < max_title_len else '...{}'.format(title[-max_title_len:])
            ))

    @staticmethod
    def get_current_event_time():
        """
        returns a valid timestamp
        :return: timestamp
        """
        return GdkX11.x11_get_server_time(Gdk.get_default_root_window())


########################################################################################################################
# Main
########################################################################################################################

if __name__ == '__main__':

    arg_parser = ArgumentParser(description='window - window manager for the command line')
    arg_parser.add_argument('-a', '--all-workspace', action='store_true',
                            help='include all windows of the active workspace')
    arg_parser.add_argument('-A', '--all', action='store_true',
                            help='include all windows')
    arg_parser.add_argument('-c', '--center', nargs='*', type=str, metavar='WIN',
                            help='center windows')
    arg_parser.add_argument('-d', '--distribute', nargs='*', type=str, metavar='WIN',
                            help='distribute windows across the screen')
    arg_parser.add_argument('-D', '--debug', action='store_true',
                            help='show debug messages')
    arg_parser.add_argument('-f', '--focus', nargs=1, type=str, metavar='WIN',
                            help='focus the given window')
    arg_parser.add_argument('-g', '--group', nargs='*', type=str, metavar='WIN',
                            help='group and cascade windows')
    arg_parser.add_argument('-G', '--geometry', action='store_true',
                            help='set geometry on opened windows')
    arg_parser.add_argument('-i', '--init', nargs='*', type=str, metavar='WIN',
                            help='initial setup of terminal + default + given WINs')
    arg_parser.add_argument('-l', '--list', action='store_true',
                            help='list windows')
    arg_parser.add_argument('--monitor', action='store_true',
                            help='monitor windows on opening')
    arg_parser.add_argument('-m', '--maximize', nargs='*', type=str, metavar='WIN',
                            help='maximize windows')
    arg_parser.add_argument('-M', '--move-workspaces', nargs='*', type=int, metavar='NR',
                            help='move workspaces; 2 args: switch two workspaces; 1 arg: insert empty workspace at ' +
                                 'the given position; 0 arg: remove gaps between workspaces')
    arg_parser.add_argument('-o', '--open', nargs='+', type=str, metavar='APP',
                            help='open windows')
    arg_parser.add_argument('-q', '--close-apps', nargs='+', type=str, metavar='APP',
                            help='close all windows of the given app')
    arg_parser.add_argument('-r', '--rearrange', action='store_true',
                            help='move windows to their configured workspaces')
    arg_parser.add_argument('-w', '--close-windows', nargs='*', type=str, metavar='WIN',
                            help='close windows')

    args = vars(arg_parser.parse_args())

    if args['debug'] or args['monitor']:
        LOG.setLevel(logging.DEBUG)
        STREAM_HANDLER.setLevel(logging.DEBUG)

    with open(path.join(path.dirname(path.abspath(__file__)), 'window.yml'), 'r') as config:
        wm = WindowManager(config)

    if args['init'] is not None:
        wm.maximize_windows(args=args)
        wm.open_apps(args['init'], default=True, set_geometry=args['geometry'])
    elif args['open']:
        wm.open_apps(args['open'], set_geometry=args['geometry'])
    elif args['monitor']:
        wm.launch_count = 1
        try:
            wm.open_apps([], default=False, set_geometry=False)
        except KeyboardInterrupt:
            print('\nmonitor stopped')
    elif args['close_windows'] is not None:
        wm.close_windows(args=args)
    elif args['close_apps'] is not None:
        wm.close_apps(args=args)
    elif args['group'] is not None:
        wm.cascade_windows(args=args)
    elif args['maximize'] is not None:
        wm.maximize_windows(args=args)
    elif args['center'] is not None:
        wm.center_windows(args=args)
    elif args['distribute'] is not None:
        wm.distribute_windows(args=args)
    elif args['focus']:
        wm.focus_window(args['focus'][0])
    elif args['rearrange']:
        wm.rearrange_windows(set_geometry=args['geometry'])
    elif args['move_workspaces'] is not None:
        wm.move_workspaces(args['move_workspaces'])
    elif args['list']:
        wm.print_window_list()
        print('')
    else:
        wm.maximize_windows(args=args)

# sample config (window.yml):

# ---
# geometry:
#   menubar: 40
#   origin: 60
#   offset: 60
#   dimx: 0.75
#   dimy: 0.75
# format:
#   spacing: 3
#   maxtitle: 80
# apps:
#   terminal:
#     alias: te
#     exec: gnome-terminal
#     workspace: 0
#     geometry: maximize
#     default: false
#   intellij:
#     alias: ij
#     exec: ~/.local/bin/idea
#     pattern: '^[^ ]+ \[[^\]]+\] - .*/.*|^welcome to intellij idea$'
#     workspace: -1
#     geometry: maximize
#     default: false
#   firefox:
#     alias: ff
#     exec: firefox
#     workspace: 1
#     geometry: maximize
#     default: true
#   evolution:
#     alias: ev
#     exec: evolution
#     workspace: 2
#     geometry: maximize
#     default: true