dotfiles/.config/yabai/yabai.py

#!/usr/bin/env python3
from os.path import exists
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 Any, 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[str, Any]):
        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[str, Any]):
        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 YabaiDisplay:
    def __init__(self, data: dict[str, Any]):
        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[str, int] = data["frame"]

    def __repr__(self) -> str:
        return (
            f"Display({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: "YabaiDisplay"):
        return self.frame["w"].__gt__(other.frame["w"])


class BetterDisplayDisplay:
    def __init__(self, data: dict[str, Any]):
        self.alphanumeric_serial: str = data["alphanumericSerial"]
        self.device_type: str = data["deviceType"]
        self.display_id: str = data["displayID"]
        self.model: str = data["model"]
        self.name: str = data["name"]
        self.original_name: str = data["originalName"]
        self.product_name: str = data["productName"]
        self.registry_location: str = data["registryLocation"]
        self.serial: str = data["serial"]
        self.tagID: str = data["tagID"]
        self.uuid: str = data["UUID"]
        self.vendor: str = data["vendor"]
        self.manufacture_week: str = data["weekOfManufacture"]
        self.manufacture_year: str = data["yearOfManufacture"]

    def __repr__(self) -> str:
        return f"BetterDisplayDisplay({self.name})"

    def __gt__(self, other: "BetterDisplayDisplay"):
        return self.display_id.__gt__(other.display_id)

    @property
    def built_in(self) -> bool:
        return self.name == "Built-in Display"


SpaceSel: TypeAlias = int | str


class CLIWrapper:
    _base_args: list[str] = []

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

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


class BetterDisplay(CLIWrapper):
    _available: bool = False
    _base_args: list[str] = [
        "/Applications/BetterDisplay.app/Contents/MacOS/BetterDisplay"
    ]

    def __init__(self):
        if exists(self._base_args[0]) and self.execute(self._base_args + ["help"], timeout=0.5):
            self._available = True

    @property
    def available(self) -> bool:
        return self._available

    def set_brightness(
        self, display: BetterDisplayDisplay, value: float
    ):
        self.message(
            ["set", f"-uuid={display.uuid}", "-feature=brightness", f"-value={value}"]
        )

    def get_displays(self) -> list[BetterDisplayDisplay]:
        return [
            BetterDisplayDisplay(display)
            for display in loads(
                "[" + self.message(["get", "-feature=identifiers"]) + "]"
            )
        ]


class Yabai(CLIWrapper):
    _base_args: list[str] = ["yabai", "-m"]
    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: Any, exc_value: Any, traceback: Any):
        if exc_type is not None:
            debug(f"Exited with {exc_type} {exc_value}")
        self.message(["rule", "--apply"])
        if self._initial_window is not None:
            self.message(["window", "--focus", self._initial_window])
        if exc_type is None:
            debug(f"Executed successfully")

    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[YabaiDisplay]:
        return {
            YabaiDisplay(display)
            for display in loads(self.message(["query", "--displays"]))
        }

    def get_main_display(
        self, displays: set[YabaiDisplay] | None = None
    ) -> YabaiDisplay:
        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"YabaiDisplay {display.index}",
                    ]
                )
        better_display = BetterDisplay()
        if better_display.available:
            displays = better_display.get_displays()
            if len(displays) > 1:
                for display in displays:
                    if display.built_in:
                        better_display.set_brightness(display, 0)
                    else:
                        better_display.set_brightness(display, 1)

    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, _ 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, _ in self.spaces:
            self.set_config(
                label,
                fullscreen=fullscreen,
                horizontal_padding=int(
                    (main_display.frame["w"] - TARGET_DISPLAY_WIDTH) / 2
                ),
                vertical_padding=int(
                    (main_display.frame["h"] - TARGET_DISPLAY_HEIGHT) / 2
                ),
            )
        if len(created_space_labels) > 0 and initial_window is not None:
            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
        if initial_window is not None:
            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",
            ]
        )

    def get_focused_window(self) -> Window | None:
        windows = [
            window for window in [Window(window) for window in loads(self.message(["query", "--windows"]))] if window.has_focus
        ]
        if len(windows) > 0:
            return windows.pop()
        return None


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