#!/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 ''}" 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 ''}" 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.id]) 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) else: display = displays.pop() if display: 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.id]) 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()