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