rescriptoon/tooncher/controls.py

import copy
import time
try:
    import psutil
except ImportError:
    psutil = False
try:
    import Xlib.display
    from Xlib import X, XK
except ImportError:
    Xlib = False

if Xlib:
    EXTENDED_KEYBOARD_CONTROLS_DEFAULT_MAPPING = {
        XK.XK_w: XK.XK_Up,
        XK.XK_a: XK.XK_Left,
        XK.XK_s: XK.XK_Down,
        XK.XK_d: XK.XK_Right,
    }


def x_find_window(parent_window, filter_callback):
    matching = []
    for child_window in parent_window.query_tree().children:
        if filter_callback(child_window):
            matching.append(child_window)
        matching += x_find_window(child_window, filter_callback)
    return matching


def x_find_window_by_pid(display, pid):
    pid_prop = display.intern_atom('_NET_WM_PID')

    def filter_callback(window):
        prop = window.get_full_property(pid_prop, X.AnyPropertyType)
        return prop and prop.value.tolist() == [pid]
    return x_find_window(display.screen().root, filter_callback)


class ExtendedControls:

    def __init__(self, engine_pid, toggle_keysym_name):
        if not psutil:
            raise Exception('\n'.join([
                'Extended keyboard controls require the python lib psutil to be installed.',
                'Depending on your system run',
                '\t$ sudo apt-get install python3-psutil',
                'or',
                '\t$ pip3 install --user psutil',
            ]))
        if not Xlib:
            raise Exception('\n'.join([
                'Extended keyboard controls require xlib for python to be installed.',
                'Depending on your system run',
                '\t$ sudo apt-get install python3-xlib',
                'or',
                '\t$ pip3 install --user xlib',
            ]))
        self._engine_pid = engine_pid
        self._xdisplay = Xlib.display.Display()
        self._toggle_keysym = XK.string_to_keysym(toggle_keysym_name)
        if self._toggle_keysym == X.NoSymbol:
            raise Exception("Extended keyboard controls toggle:"
                            + " Unknown keysym name '{}'".format(toggle_keysym_name))
        keyboard_mapping = copy.deepcopy(
            EXTENDED_KEYBOARD_CONTROLS_DEFAULT_MAPPING
        )
        if self._toggle_keysym in keyboard_mapping:
            del keyboard_mapping[self._toggle_keysym]
            print("INFO Extended Controls:"
                  + " Ignoring mapping for toggle key '{}'".format(toggle_keysym_name))
        self._keyboard_mapping = keyboard_mapping
        self._engine_window = None
        self._enabled = False

    @property
    def engine_running(self):
        return psutil.pid_exists(self._engine_pid)

    def _wait_for_engine_window(self, timeout_seconds=20, search_interval_seconds=2):
        start_epoch = time.time()
        while self.engine_running and (time.time() - start_epoch) <= timeout_seconds:
            windows = x_find_window_by_pid(
                self._xdisplay,
                self._engine_pid,
            )
            assert len(windows) <= 1
            if len(windows) == 1:
                return windows[0]
            time.sleep(search_interval_seconds)
        return None

    def run(self):
        self._engine_window = self._wait_for_engine_window()
        if not self._engine_window:
            raise Exception('Could not find the game\'s window.')
        self._grab_key(
            self._xdisplay.keysym_to_keycode(self._toggle_keysym),
        )
        if not self.enabled:
            keysym_name = XK.keysym_to_string(self._toggle_keysym)
            print("INFO Extended Controls are currently disabled."
                  + " Press key '{}' to enable.".format(keysym_name))
        while self.engine_running:
            # TODO don't block here, engine might have already been stopped
            self._handle_xevent(self._xdisplay.next_event())

    def _handle_xevent(self, xevent):
        # TODO investigate why some release events get lost
        if isinstance(xevent, Xlib.protocol.event.KeyPress) \
                or isinstance(xevent, Xlib.protocol.event.KeyRelease):
            self._handle_xkeyevent(xevent)

    def _handle_xkeyevent(self, xkeyevent):
        # TODO map keycodes instead of keysyms
        keysym_in = self._xdisplay.keycode_to_keysym(
            xkeyevent.detail,
            index=0,
        )
        if keysym_in == self._toggle_keysym:
            if isinstance(xkeyevent, Xlib.protocol.event.KeyPress):
                self.toggle()
        else:
            if self.enabled and keysym_in in self._keyboard_mapping:
                keysym_out = self._keyboard_mapping[keysym_in]
            else:
                keysym_out = keysym_in
            self._engine_window.send_event(type(xkeyevent)(
                window=self._engine_window,
                detail=self._xdisplay.keysym_to_keycode(keysym_out),
                state=0,
                root_x=xkeyevent.root_x,
                root_y=xkeyevent.root_y,
                event_x=xkeyevent.event_x,
                event_y=xkeyevent.event_y,
                child=xkeyevent.child,
                root=xkeyevent.root,
                time=xkeyevent.time,  # X.CurrentTime
                same_screen=xkeyevent.same_screen,
            ))

    def enable(self):
        for keysym in self._keyboard_mapping.keys():
            self._grab_key(
                self._xdisplay.keysym_to_keycode(keysym),
            )
        self._enabled = True
        print("INFO Enabled Extended Controls")

    def disable(self):
        for keysym in self._keyboard_mapping.keys():
            self._ungrab_key(
                self._xdisplay.keysym_to_keycode(keysym),
            )
        self._enabled = False
        print("INFO Disabled Extended Controls")

    @property
    def enabled(self):
        return self._enabled

    def toggle(self):
        if self.enabled:
            self.disable()
        else:
            self.enable()

    def _grab_key(self, keycode):
        self._engine_window.grab_key(
            keycode,
            X.AnyModifier,
            # owner_events
            # https://stackoverflow.com/questions/32122360/x11-will-xgrabpointer-prevent-other-apps-from-any-mouse-event
            # False,
            True,
            X.GrabModeAsync,
            X.GrabModeAsync,
        )

    def _ungrab_key(self, keycode):
        self._engine_window.ungrab_key(keycode, X.AnyModifier)