dotfiles-sway/sway/scripts/per-window-layout.py

#!/usr/bin/env python3

from __future__ import annotations

import asyncio
import logging
import os
import signal
from dataclasses import dataclass, field
from typing import Dict, Optional

try:
    from systemd import daemon
except ImportError:
    daemon = None

import i3ipc
from i3ipc.aio import Connection

DEFAULT_MODE = os.getenv("PWL_DEFAULT_MODE", "1") == "1"
DEBUG = os.getenv("PWL_DEBUG", "0") == "1"
INFO = os.getenv("PWL_INFO", "1") == "1"

logging.basicConfig(
    level=logging.DEBUG if DEBUG else (logging.INFO if INFO else logging.WARNING),
    format="[pwl] %(asctime)s %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("pwl")


def notify_systemd(msg: str) -> None:
    if daemon:
        try:
            daemon.notify(msg)
        except Exception as e:
            log.debug("Failed to notify systemd: %s", e)


def notify_ready(status: str) -> None:
    notify_systemd(f"READY=1\nSTATUS={status}")


def notify_status(status: str) -> None:
    notify_systemd(f"STATUS={status}")


def notify_watchdog() -> None:
    notify_systemd("WATCHDOG=1")


@dataclass
class Keyboard:
    identifier: str
    default_index: int
    current_index: int
    layout_map: Dict[str, int] = field(default_factory=dict)


class PerWindowLayout:
    def __init__(self) -> None:
        self.conn: Optional[Connection] = None
        self.keyboard: Optional[Keyboard] = None

        self.window_layouts: Dict[int, int] = {}
        self.current_window: Optional[int] = None

        self.stop_event = asyncio.Event()
        self._layout_map_built = False

    async def connect(self) -> None:
        self.conn = await Connection(auto_reconnect=True).connect()

    def _find_main_keyboard(self, inputs) -> Optional[object]:
        keyboards = [i for i in inputs if i.type == "keyboard"]
        if not keyboards:
            return None

        real = [k for k in keyboards if "virtual" not in (k.identifier or "").lower()]
        return (real or keyboards)[0]

    async def wait_until_ready(self) -> None:
        notify_status("Waiting for sway...")
        while not self.stop_event.is_set():
            try:
                inputs = await self.conn.get_inputs()
            except Exception:
                await asyncio.sleep(0.2)
                continue

            main = self._find_main_keyboard(inputs)
            if main is None:
                await asyncio.sleep(0.2)
                continue

            idx = main.xkb_active_layout_index or 0
            self.keyboard = Keyboard(
                identifier=main.identifier,
                default_index=idx,
                current_index=idx,
            )
            return

    async def build_layout_map(self, force: bool = False) -> None:
        assert self.keyboard is not None
        inputs = await self.conn.get_inputs()
        current = self.keyboard.current_index

        kbd = next((i for i in inputs if i.identifier == self.keyboard.identifier), None)
        if not kbd:
            return

        layouts = kbd.xkb_layout_names or []
        count = len(layouts)

        if self._layout_map_built and not force and len(self.keyboard.layout_map) == count:
            return

        notify_status("Building layout map...")
        self.keyboard.layout_map.clear()

        for i in range(count):
            await self.conn.command(f'input "{self.keyboard.identifier}" xkb_switch_layout {i}')
            inputs = await self.conn.get_inputs()
            kbd = next((k for k in inputs if k.identifier == self.keyboard.identifier), kbd)
            name = kbd.xkb_active_layout_name or str(i)
            self.keyboard.layout_map[name] = i

        await self.conn.command(f'input "{self.keyboard.identifier}" xkb_switch_layout {current}')
        self.keyboard.current_index = current
        self._layout_map_built = True
        log.info("Layout map initialized (%d layouts)", len(self.keyboard.layout_map))

    async def init_windows(self) -> None:
        assert self.keyboard is not None
        inputs = await self.conn.get_inputs()
        kbd = next((i for i in inputs if i.identifier == self.keyboard.identifier), None)
        idx = kbd.xkb_active_layout_index if kbd else self.keyboard.default_index

        tree = await self.conn.get_tree()
        active_ids = set()

        for con in tree.leaves():
            active_ids.add(con.id)
            if con.id not in self.window_layouts:
                self.window_layouts[con.id] = idx

        stale = [wid for wid in self.window_layouts if wid not in active_ids]
        for wid in stale:
            self.window_layouts.pop(wid, None)

        notify_status(f"Tracking {len(self.window_layouts)} windows")

    async def switch_layout(self, idx: int) -> None:
        assert self.keyboard is not None
        if idx == self.keyboard.current_index:
            return
        await self.conn.command(f'input "{self.keyboard.identifier}" xkb_switch_layout {idx}')
        self.keyboard.current_index = idx

    async def handle_active_window(self, con_id: int) -> None:
        assert self.keyboard is not None
        self.current_window = con_id

        if con_id in self.window_layouts:
            idx = self.window_layouts[con_id]
        else:
            idx = self.keyboard.default_index if DEFAULT_MODE else self.keyboard.current_index
            self.window_layouts[con_id] = idx
            notify_status(f"Tracking {len(self.window_layouts)} windows")

        await self.switch_layout(idx)

    def handle_active_layout(self, identifier: str, layout_name: Optional[str]) -> None:
        assert self.keyboard is not None
        if identifier != self.keyboard.identifier or not layout_name:
            return

        idx = self.keyboard.layout_map.get(layout_name)
        if idx is None:
            return

        self.keyboard.current_index = idx
        if self.current_window is not None:
            self.window_layouts[self.current_window] = idx

    def handle_close_window(self, con_id: int) -> None:
        if self.window_layouts.pop(con_id, None) is not None:
            notify_status(f"Tracking {len(self.window_layouts)} windows")

    async def handle_device_change(self) -> None:
        inputs = await self.conn.get_inputs()
        main = self._find_main_keyboard(inputs)
        if main and self.keyboard and main.identifier != self.keyboard.identifier:
            log.info("Main keyboard changed: %s -> %s", self.keyboard.identifier, main.identifier)
            idx = main.xkb_active_layout_index or 0
            self.keyboard.identifier = main.identifier
            self.keyboard.default_index = idx
            self.keyboard.current_index = idx
            await self.build_layout_map(force=True)
            await self.init_windows()

    def _register_handlers(self) -> None:
        assert self.conn is not None

        async def on_window(_conn, event):
            change = event.change
            if change == "focus" and event.container is not None:
                await self.handle_active_window(event.container.id)
            elif change == "close" and event.container is not None:
                self.handle_close_window(event.container.id)

        async def on_input(_conn, event):
            change = event.change
            if change == "xkb_layout" and event.input is not None:
                self.handle_active_layout(
                    event.input.identifier, event.input.xkb_active_layout_name
                )
            elif change in ("added", "removed"):
                await self.handle_device_change()

        self.conn.on(i3ipc.Event.WINDOW, on_window)
        self.conn.on(i3ipc.Event.INPUT, on_input)

    async def watchdog_loop(self) -> None:
        interval = 10.0
        if daemon:
            try:
                wd = daemon.watchdog_enabled()
                if isinstance(wd, (int, float)) and wd > 0:
                    interval = wd / 2.0
            except Exception:
                pass

        while not self.stop_event.is_set():
            notify_watchdog()
            await asyncio.sleep(interval)

    async def run(self) -> None:
        await self.connect()
        await self.wait_until_ready()
        await self.build_layout_map()
        await self.init_windows()
        self._register_handlers()

        notify_ready(f"Tracking {len(self.window_layouts)} windows")

        async with asyncio.TaskGroup() as tg:
            tg.create_task(self.conn.main())
            tg.create_task(self.watchdog_loop())
            await self.stop_event.wait()
            await self.conn.main_quit()


async def amain() -> None:
    pwl = PerWindowLayout()
    loop = asyncio.get_running_loop()

    def shutdown_handler() -> None:
        pwl.stop_event.set()

    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, shutdown_handler)

    try:
        await pwl.run()
    except asyncio.CancelledError:
        raise
    except Exception:
        log.exception("Fatal error")
        raise


def main() -> None:
    try:
        asyncio.run(amain())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()