dotfiles/.config/yabai/yabai.py

#!/usr/bin/env python3

# pyright: strict, reportAny=false, reportExplicitAny=false, reportUnusedCallResult=false

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, Literal, TypeAlias, cast, override

HOME = getenv("HOME")
XDG_CONFIG_HOME = getenv("XDG_CONFIG_HOME")
TARGET_DISPLAY_WIDTH = 1920
TARGET_DISPLAY_HEIGHT = 1080


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

    @override
    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.display: int = data["display"]
        self.windows: list[int] = data["label"]
        self.is_visible: str = data["is-visible"]
        self.has_focus: bool = data["has-focus"]
        self.is_native_fullscreen: bool = data["is-native-fullscreen"]

    @override
    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", {self.display}"
            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"]

    @override
    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"])


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 cast(str, check_output(*args, **kwargs))


class Yabai(CLIWrapper):
    _base_args: list[str] = ["yabai", "-m"]

    # Label, Full-Page, Apps, Main Display
    spaces: list[tuple[str, bool, list[str], bool]] = [
        ("Desktop", False, ["*"], True),
        ("Finder", False, ["Finder"], True),
        ("Terminal", True, ["Alacritty"], True),
        ("Browser", True, ["Firefox"], True),
    ]
    _initial_window: Window | None = None
    _initial_spaces: list[Space] = []
    _dual_display: None | Literal[True] | Literal[False] = None
    _exit_with_rule_apply: bool = False
    _exit_with_refocus: bool = False
    _display_labels: tuple[str, str] = ("Main", "Secondary")

    def __init__(self):
        self._dual_display = len(self.get_displays()) > 1
        self.spaces.append(
            (
                "Communication",
                False,
                ["Slack", "Signal", "Spotify"],
                False,
            )
        )
        self.spaces.append(("Notetaking", True, ["Obsidian", "Asana", "Notion"], True))
        self.spaces.append(("Admin", True, ["Notion Calendar", "Notion Mail"], True))
        if self._dual_display:
            for application in [
                "Google Meet",
                "Netflix",
                "YouTube",
            ]:
                self.spaces.append(
                    (
                        application,
                        True,
                        [application],
                        False,
                    )
                )

    def __enter__(self):
        self._initial_window = self.get_focused_window()
        self._initial_spaces = [
            s for s in self.get_spaces() if s.is_visible and len(s.label)
        ]
        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}")
        if self._exit_with_rule_apply:
            self.message(["rule", "--apply"])
        if self._exit_with_refocus or self._exit_with_rule_apply:
            for space in self._initial_spaces:
                try:
                    self.message(["space", "--focus", space.label])
                except CalledProcessError:
                    pass
            if self._initial_window is not None:
                self.message(["window", "--focus", self._initial_window.id])
        if exc_type is None:
            debug("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 is not 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)
        secondary_display = None
        for display in displays:
            if display.index == main_display.index:
                self.message(
                    ["display", display.index, "--label", self._display_labels[0]]
                )
            elif secondary_display is None:
                self.message(
                    ["display", display.index, "--label", self._display_labels[1]]
                )
                secondary_display = display
            else:
                self.message(
                    [
                        "display",
                        display.index,
                        "--label",
                        f"YabaiDisplay {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, _, _ 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.id])
        # Return focus
        if initial_window is not None:
            self.message(["window", "--focus", initial_window.id])
        if self._dual_display:
            display_by_space_label = {s[0]: 1 if s[-1] else 2 for s in self.spaces}
            spaces = self.get_spaces()
            wrong_spaces = [
                s
                for s in spaces
                if s.label in display_by_space_label
                and s.display != display_by_space_label[s.label]
            ]
            wrong_main_spaces = [s for s in wrong_spaces if s.display != 1]
            wrong_secondary_spaces = [s for s in wrong_spaces if s.display != 2]
            while len(wrong_main_spaces) and len(wrong_secondary_spaces):
                self.message(
                    [
                        "space",
                        wrong_main_spaces.pop().label,
                        "--swap",
                        wrong_secondary_spaces.pop().label,
                    ]
                )
            wrong_spaces = wrong_main_spaces + wrong_secondary_spaces
            while len(wrong_spaces) and (space := wrong_spaces.pop().label):
                self.message(
                    ["space", space, "--display", display_by_space_label[space]]
                )
        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_global_config(
        self,
    ):
        self.message(["config", "auto_balance", "on"])
        self.message(["config", "mouse_follows_focus", "on"])

    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}",
                    ]
                )

        # 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",
            ]
        )
        # Run Slack client script (populating cached values) when Slack is launched
        self.message(
            [
                "signal",
                "--add",
                "event=application_launched",
                "app=Slack",
                "label=SlackRunSlackClientScript",
                f"action=/usr/bin/env python3 {HOME}/.scripts/slack_client.py",
            ]
        )
        # Check if dark mode settings have been updated when focusing terminal
        self.message(
            [
                "signal",
                "--add",
                "event=window_focused",
                "app=Alacritty",
                "label=AlacrittyCheckDarkMode",
                "action=/bin/zsh $DOTFILES_DIR/.scripts/lightmode.zsh",
            ]
        )

        if self._dual_display:
            self.message(
                [
                    "signal",
                    "--add",
                    "event=window_focused",
                    "label=DisplayBrightnessManager",
                    "action=/bin/zsh $DOTFILES_DIR/.scripts/display_brightness.zsh",
                ]
            )

        for trigger in [
            "display_added",
            "display_removed",
            "display_resized",
            "system_woke",
        ]:
            self.message(
                [
                    "signal",
                    "--add",
                    f"event={trigger}",
                    f"label=DisplayChangeYabai{trigger}",
                    f"action=/usr/bin/env python3 {HOME}/.config/yabai.py manage",
                ]
            )
            self.message(
                [
                    "signal",
                    "--add",
                    f"event={trigger}",
                    f"label=DisplayChangeDisplay{trigger}",
                    "action=/bin/zsh $DOTFILES_DIR/.scripts/display_brightness.zsh",
                ]
            )

        # Rules that differ for one or multiple displays
        if not self._dual_display:
            # 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",
                    ]
                )
            # Tiny streaming player
            self.message(
                [
                    "rule",
                    "--add",
                    "label=NetflixFloatingWindowRule",
                    "app=Netflix",
                    "sticky=on",
                    "manage=on",
                ]
            )
        else:
            self.message(
                [
                    "rule",
                    "--add",
                    "label=SlackVideoCallFloatingWindowRule",
                    "app=Slack",
                    "title=Huddle.*",
                    "space=Google Meet",
                ]
            )

    def move_spaces_to_displays(self):
        if not self._dual_display:
            return
        current_spaces = self.get_spaces()
        displays = self.get_displays()
        main_display = next(
            (d for d in displays if d.label == self._display_labels[0]), None
        )
        secondary_display = next(
            (d for d in displays if d.label == self._display_labels[1]), None
        )
        if main_display is None or secondary_display is None:
            return

        incorrect_main_spaces: set[Space] = set()
        incorrect_secondary_spaces: set[Space] = set()
        for space_label, _, _, is_main_display_space in self.spaces:
            mapped_space = next(
                (s for s in current_spaces if s.label == space_label), None
            )
            if mapped_space is None:
                continue

            if (
                mapped_space.display == main_display.index
                and not is_main_display_space
                and secondary_display
            ):
                incorrect_secondary_spaces.add(mapped_space)
            elif mapped_space.display != main_display.index and is_main_display_space:
                incorrect_main_spaces.add(mapped_space)
        last_focus = next((s.label for s in current_spaces if s.has_focus), None)
        while len(incorrect_main_spaces) > 0 and len(incorrect_secondary_spaces) > 0:
            from_space = incorrect_main_spaces.pop()
            to_space = incorrect_secondary_spaces.pop()
            if to_space.label == last_focus:
                from_space, to_space = to_space, from_space
            self.message(
                [
                    "space",
                    from_space.label,
                    "--switch",
                    to_space.label,
                ]
            )
            last_focus = to_space.label
        for spaces, secondary in (
            (incorrect_main_spaces, False),
            (incorrect_secondary_spaces, True),
        ):
            for space in spaces:
                self.message(
                    [
                        "space",
                        space.label,
                        "--display",
                        secondary_display.label
                        if secondary and secondary_display
                        else main_display.label,
                    ]
                )

    def sort_spaces(self):
        all_spaces = [s.label for s in sorted(self.get_spaces())]
        for is_main in (True, False) if self._dual_display else [None]:
            order = [
                s[0]
                for s in self.spaces
                if s[0] in all_spaces and (s[3] == is_main) or is_main is None
            ]
            spaces = [s for s in all_spaces if s in order]

            while tuple(spaces) != tuple(order) and len(spaces) == len(order):
                for index in range(1, len(spaces)):
                    if order.index(spaces[index]) < order.index(spaces[index - 1]):
                        self.message(["space", spaces[index], "--move", "prev"])
                        spaces[index], spaces[index - 1] = (
                            spaces[index - 1],
                            spaces[index],
                        )

    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

    def enable_exit_with_rule_apply(self):
        self._exit_with_rule_apply = True
        self.enable_exit_with_refocus()

    def enable_exit_with_refocus(self):
        self._exit_with_refocus = True


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.set_global_config()
            yabai.enable_exit_with_rule_apply()
            yabai.manage_displays()
            yabai.manage_spaces()
        if argv[1] == "initialize":
            yabai.set_rules_and_signals()
        if argv[1] == "move" or argv[1] == "manage" or argv[1] == "initialize":
            yabai.enable_exit_with_refocus()
            yabai.move_spaces_to_displays()