dotfiles/.config/yabai/yabai.py

#!/usr/bin/env python3
from json import loads
from logging import NOTSET, basicConfig, debug, info
from os import getenv
from subprocess import CalledProcessError, check_output
from sys import argv
from typing import TypeAlias

HOME = getenv("HOME")
XDG_CONFIG_HOME = getenv("XDG_CONFIG_HOME")
TARGET_DISPLAY_WIDTH = 1680
TARGET_DISPLAY_HEIGHT = 1050

class Window:
    def __init__(self, data: dict):
        self.id: int = data["id"]
        self.app: str = data["app"]
        self.space: str = data["space"]
        self.title: str = data["title"]
        self.has_focus: bool = data["has-focus"]
        self.frame_w: int = data["frame"]["w"]
        self.frame_h: int = data["frame"]["h"]

    @property
    def size(self) -> int:
        return self.frame_w * self.frame_h

    def __repr__(self) -> str:
        return (
            f"Window({self.title} | {self.app}"
            f", {self.id}"
            f"{', Focused' if self.has_focus else ''})"
        )

    def __gt__(self, other: "Window"):
        return self.id.__gt__(other.id)
 

class Space:
    def __init__(self, data: dict):
        self.id: int = data["id"]
        self.index: int = data["index"]
        self.label: str = data["label"]
        self.windows: list[int] = data["label"]
        self.has_focus: bool = data["has-focus"]
        self.is_native_fullscreen: bool = data["is-native-fullscreen"]

    def __repr__(self) -> str:
        return (
            f"Space({self.label if self.label and len(self.label) > 0 else '<NoLabel>'}"
            f", {self.index}"
            f"{', Fullscreen' if self.is_native_fullscreen else ''}"
            f"{', Focused' if self.has_focus else ''})"
        )

    def __gt__(self, other: "Space"):
        return self.index.__gt__(other.index)


class Display:
    def __init__(self, data: dict):
        self.id: int = data["id"]
        self.uuid: str = data["uuid"]
        self.index: int = data["index"]
        self.label: str = data["label"]
        self.has_focus: bool = data["has-focus"]
        self.spaces: list[str] = data["spaces"]
        self.frame: dict = data["frame"]

    def __repr__(self) -> str:
        return (
            f"Space({self.label if self.label and len(self.label) > 0 else '<NoLabel>'}"
            f", {self.index}"
            f", Frame ({self.frame})"
            f"{', Focused' if self.has_focus else ''})"
        )

    def __gt__(self, other: "Display"):
        return self.frame["w"].__gt__(other.frame["w"])


SpaceSel: TypeAlias = int | str


class Yabai:
    spaces: list[tuple[str, bool, list[str]]] = [
        ("Desktop", False, ["*"]),
        ("Finder", False, ["Finder"]),
        ("Terminal", True, ["Alacritty"]),
        ("Browser", True, ["Firefox"]),
        ("Communication", False, ["Slack", "Signal", "Spotify", "Notion"]),
        ("Notetaking", True, ["Obsidian", "Asana"]),
    ]
    _initial_window: int | None = None

    def __enter__(self):
        self._initial_window = self.get_focused_window()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            debug(f"Exited with {exc_type} {exc_value}")
        self.message(["rule", "--apply"])
        self.message(["window", "--focus", self._initial_window])
        if exc_type is None:
            debug(f"Executed successfully")

    def execute(self, *args, **kwargs):
        debug(f"Executing: ({args}), ({kwargs})")
        return check_output(*args, **kwargs)

    def message(self, args: list[str | int]) -> str:
        return self.execute(["yabai", "-m"] + [str(arg) for arg in args]).decode("utf8")

    def get_windows(self) -> set[Window]:
        return { Window(window) for window in loads(self.message(["query", "--windows"])) }

    def get_spaces(self) -> set[Space]:
        return {Space(space) for space in loads(self.message(["query", "--spaces"]))}

    def get_displays(self) -> set[Display]:
        return { Display(display) for display in loads(self.message(["query", "--displays"])) }

    def get_main_display(self, displays: set[Display] | None = None) -> Display:
        return sorted(list(displays if displays != None else self.get_displays()))[-1]

    def is_blank_space(self, space: Space) -> bool:
        return (
            space.label not in {s[0] for s in self.spaces}
            and not space.is_native_fullscreen
        )

    def manage_displays(self):
        displays = self.get_displays()
        main_display = self.get_main_display(displays)
        for display in displays:
            if display.index == main_display.index:
                self.message(["display", display.index, "--label", "Main"])
            else:
                self.message(
                    ["display", display.index, "--label", f"Display {display.index}"]
                )

    def manage_spaces(self):
        initial_window = self.get_focused_window()
        # Start by making sure that the expected number of spaces are present
        spaces = self.get_spaces()
        occupied_spaces = {
            space
            for space in spaces
            if not self.is_blank_space(space) and not space.is_native_fullscreen
        }
        blank_spaces = {space for space in spaces if self.is_blank_space(space)}
        for _ in range(
            max(0, len(self.spaces) - (len(occupied_spaces) + len(blank_spaces)))
        ):
            self.message(["space", "--create"])
        # Use blank spaces to create occupied spaces as necessary
        spaces = self.get_spaces()
        blank_spaces = {space for space in spaces if self.is_blank_space(space)}
        created_space_labels: set[tuple[str, bool]] = set()
        for space_label, space_fullscreen, space_apps in self.spaces:
            if any(space.label == space_label for space in spaces):
                continue
            space = blank_spaces.pop()
            self.message(["space", space.index, "--label", space_label])
            created_space_labels.add((space_label, space_fullscreen))
        # Remove unnecessary spaces
        spaces = self.get_spaces()
        blank_spaces = [space for space in spaces if self.is_blank_space(space)]
        blank_spaces.sort(key=lambda s: s.index, reverse=True)
        # Make sure that the focused space isn't a blank one
        if any(space.has_focus for space in blank_spaces):
            self.message(["space", "--focus", self.spaces[0][0]])
        for space in blank_spaces:
            self.message(["space", "--destroy", space.index])
        # Configure the new spaces
        main_display = self.get_main_display()
        for label, fullscreen in created_space_labels:
            self.set_space_background(label)
        for label, fullscreen, apps in self.spaces:
            self.set_config(
                label,
                fullscreen=fullscreen,
                horizontal_padding=(main_display.frame["w"] - TARGET_DISPLAY_WIDTH) / 2,
                vertical_padding=(main_display.frame["h"] - TARGET_DISPLAY_HEIGHT) / 2,
            )
        if len(created_space_labels) > 0:
            self.message(["window", "--focus", initial_window])
        # Sort the remaining spaces
        for space_index, space_label in enumerate(self.spaces):
            try:
                self.message(["space", space_label[0], "--move", space_index + 1])
            except CalledProcessError:
                # Almost certainly thrown because space is already in place, so no problem
                pass
        # Return focus
        self.message(["window", "--focus", initial_window])
        info(f"Spaces configured: {sorted(self.get_spaces())}")

    def set_space_background(self, space: SpaceSel):
        try:
            self.message(["space", "--focus", space])
        except CalledProcessError:
            # Almost certainly thrown because space is already focused, so no problem
            pass
        self.execute(
            [
                "osascript",
                "-e",
                'tell application "System Events" to tell every desktop to set picture to "/System/Library/Desktop Pictures/Solid Colors/Black.png"',
            ]
        )

    def set_config(
        self,
        space: SpaceSel,
        fullscreen: bool = False,
        gap: int = 10,
        vertical_padding: int = 0,
        horizontal_padding: int = 0,
    ):
        for config in [
            ["window_shadow", "float"],
            ["window_opacity", "on"],
            ["layout", "bsp"],
            ["top_padding", int(max(0 if fullscreen else gap, vertical_padding))],
            ["bottom_padding", int(max(0 if fullscreen else gap, vertical_padding))],
            ["left_padding", int(max(0 if fullscreen else gap, horizontal_padding))],
            ["right_padding", int(max(0 if fullscreen else gap, horizontal_padding))],
            ["window_gap", gap],
        ]:
            self.message(["config", "--space", space] + config)

    def set_rules_and_signals(self):
        # Reset rules and signals
        for domain in ["rule", "signal"]:
            for _ in range(len(loads(self.message([domain, "--list"])))):
                self.message([domain, "--remove", 0])
        # Load the system agent on dock restart
        self.message(
            [
                "signal",
                "--add",
                "label=SystemAgentReloadSignal",
                "event=dock_did_restart",
                "action=sudo yabai --load-sa",
            ]
        )
        # Reload spaces when displays are reset
        for reload_event in ["display_added", "display_removed"]:
            self.message(
                [
                    "signal",
                    "--add",
                    f"label={reload_event}RestartSignal",
                    f"event={reload_event}",
                    f"action=/bin/zsh {XDG_CONFIG_HOME}/yabai/yabairc",
                ]
            )
        # Normal windows should be put on the desktop
        self.message(
            [
                "rule",
                "--add",
                "label=DefaultDesktopRule",
                "subrole=AXStandardWindow",
                "space=^Desktop",
            ]
        )
        # Rules for applications that get their own spaces
        for label, _, apps in self.spaces:
            for app in apps:
                if app == "*":
                    continue
                self.message(
                    [
                        "rule",
                        "--add",
                        f"label={app}{label}Rule",
                        f"app={app}",
                        f"space=^{label}",
                    ]
                )
        # Google Meet and Slack Huddles should be "sticky"
        for app, title in (("Google Meet", ".*"), ("Slack", "Huddle.*")):
            self.message(
                [
                    "rule",
                    "--add",
                    f"label={app}VideoCallFloatingWindowRule",
                    f"app={app}",
                    f"title={title}",
                    "sticky=on",
                    "manage=on",
                    "opacity=0.9",
                    "grid=10:10:6:6:4:4",
                ]
            )
        # Compile SurfingKeys configuration when Firefox is launched
        self.message(
            [
                "signal",
                "--add",
                "event=application_launched",
                "app=Firefox",
                "label=FirefoxCompileExtensionsSignal",
                f"action=/bin/zsh {XDG_CONFIG_HOME}/surfingkeys/compile.sh",
            ]
        )
        # Check if dark mode settings have been updated when focusing terminal
        self.message(
            [
                "signal",
                "--add",
                "event=window_focused",
                "app=Alacritty",
                "label=AlacrittyCheckDarkMode",
                f"action=/bin/zsh {HOME}/.scripts/lightmode.zsh",
            ]
        )
        # When focusing an that may share a space with at least two others, rotate the
        # space's windows to ensure that the focused window is the largest.
        for label, _, apps in self.spaces:
            if len(apps) < 3:
                continue
            for app in apps:
                self.message(
                [
                    "signal",
                    "--add",
                    "event=window_focused",
                    f"app={app}",
                    f"label={app}RotateToFocus",
                    f"action=python3 {XDG_CONFIG_HOME}/yabai/yabai.py rotate",
                ]
            )

    def rotate(self) -> None:
        windows = self.get_windows()
        focus_window = next(w for w in windows if w.has_focus)
        same_space_windows = sorted((w for w in windows if w.space == focus_window.space), key = lambda w: w.size)
        if same_space_windows[-1].has_focus:
            return
        self.message(["window", focus_window.id, "--swap", same_space_windows[-1].id])
        return

    def get_focused_window(self) -> int:
        return [
            window
            for window in loads(self.message(["query", "--windows"]))
            if window["has-focus"]
        ].pop()["id"]


if __name__ == "__main__":
    basicConfig(level=NOTSET)
    debug(f"Called with parameters {argv}")
    if argv[1] == "manage" or argv[1] == "initialize":
        with Yabai() as yabai:
            yabai.manage_displays()
            yabai.manage_spaces()
        if argv[1] == "initialize":
            yabai.set_rules_and_signals()
    if argv[1] == "rotate":
        Yabai().rotate()