#!/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()