From a6a565ec60c955b952666bf1ac90b8cfd0451e98 Mon Sep 17 00:00:00 2001 From: Dillon DuPont Date: Mon, 12 May 2025 22:08:42 -0400 Subject: [PATCH] added dock and menubar bounds --- notebooks/app_screenshots.py | 63 +++++------ notebooks/ui_bounds_helper.py | 199 ++++++++++++++++++++++++++++++++++ 2 files changed, 231 insertions(+), 31 deletions(-) create mode 100644 notebooks/ui_bounds_helper.py diff --git a/notebooks/app_screenshots.py b/notebooks/app_screenshots.py index 36c6ca4b..f10424c1 100644 --- a/notebooks/app_screenshots.py +++ b/notebooks/app_screenshots.py @@ -16,6 +16,11 @@ import io import asyncio import functools +from ui_bounds_helper import ( + get_menubar_bounds, + get_dock_bounds, +) + # Timing decorator for profiling def timing_decorator(func): @functools.wraps(func) @@ -256,15 +261,12 @@ class AppScreenshotCapture: z_index = z_order.get(window_id, -1) # Determine window role (desktop, dock, menubar, app) - if window_owner == "Window Server": - if window_name == "Desktop": - role = "desktop" - elif window_name == "Dock": - role = "dock" - elif window_name == "Menubar": - role = "menubar" - elif window_owner == "Dock": + if window_name == "Dock" and window_owner == "Dock": role = "dock" + elif window_name == "Menubar" and window_owner == "Window Server": + role = "menubar" + elif window_owner in ["Window Server", "Dock"]: + role = "desktop" else: role = "app" @@ -319,25 +321,15 @@ class AppScreenshotCapture: if all_windows is None: all_windows = self.get_all_windows() - # Create a list of window IDs to include - window_ids_to_include = [] - # Reverse the list to get the correct z-order - for window in all_windows[::-1]: - owner = window["owner"] - name = window["name"] - role = window["role"] - is_on_screen = window["is_on_screen"] - - # Skip windows that are not on screen - if not is_on_screen: - continue - - # Skip app windows that are not in the filter - if app_filter is not None and owner not in app_filter and role == "app": - continue - - window_ids_to_include.append(window["id"]) + all_windows = all_windows[::-1] + + # Filter out windows that are not on screen + all_windows = [window for window in all_windows if window["is_on_screen"]] + + # Filter out windows that are not in the app_filter + if app_filter is not None: + all_windows = [window for window in all_windows if window["owner"] in app_filter or window["role"] != "app"] # Get the main screen dimensions main_screen = AppKit.NSScreen.mainScreen() @@ -363,9 +355,9 @@ class AppScreenshotCapture: ) else: # Create a CFArray of window IDs to include - window_list = Foundation.CFArrayCreateMutable(None, len(window_ids_to_include), None) - for window_id in window_ids_to_include: - Foundation.CFArrayAppendValue(window_list, window_id) + window_list = Foundation.CFArrayCreateMutable(None, len(all_windows), None) + for window in all_windows: + Foundation.CFArrayAppendValue(window_list, window["id"]) # Capture only the specified windows cg_image = Quartz.CGWindowListCreateImageFromArray( @@ -387,6 +379,7 @@ class AppScreenshotCapture: image_data = io.BytesIO(png_data) return Image.open(image_data) + @timing_decorator def get_menubar_items(self, active_app_pid: int = None) -> List[Dict[str, Any]]: """Get menubar items from the active application using Accessibility API @@ -662,6 +655,9 @@ class AppScreenshotCapture: result["applications"].append(app_data) + # Add all windows to the result + result["windows"] = all_windows + # Get menubar items from the active application menubar_items = self.get_menubar_items(active_app_pid) result["menubar_items"] = menubar_items @@ -670,8 +666,13 @@ class AppScreenshotCapture: dock_items = self.get_dock_items() result["dock_items"] = dock_items - # Add all windows to the result - result["windows"] = all_windows + # Get menubar bounds + menubar_bounds = get_menubar_bounds() + result["menubar_bounds"] = menubar_bounds + + # Get dock bounds + dock_bounds = get_dock_bounds() + result["dock_bounds"] = dock_bounds # Capture the entire desktop using Quartz compositing desktop_screenshot = self.capture_desktop_screenshot(app_filter, all_windows) diff --git a/notebooks/ui_bounds_helper.py b/notebooks/ui_bounds_helper.py new file mode 100644 index 00000000..da5d0ec8 --- /dev/null +++ b/notebooks/ui_bounds_helper.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +""" +UI Bounds Helper - A utility to get accurate bounds for macOS UI elements + +This module provides helper functions to get accurate bounds for macOS UI elements +like the menubar and dock, which are needed for proper screenshot composition. +""" + +import sys +import time +from typing import Dict, Any, Optional, Tuple + +# Import Objective-C bridge libraries +try: + import AppKit + from ApplicationServices import ( + AXUIElementCreateSystemWide, + AXUIElementCreateApplication, + AXUIElementCopyAttributeValue, + AXUIElementCopyAttributeValues, + kAXChildrenAttribute, + kAXRoleAttribute, + kAXTitleAttribute, + kAXPositionAttribute, + kAXSizeAttribute, + kAXErrorSuccess, + AXValueGetType, + kAXValueCGSizeType, + kAXValueCGPointType, + AXUIElementGetTypeID, + AXValueGetValue, + kAXMenuBarAttribute, + ) + from AppKit import NSWorkspace, NSRunningApplication + import Foundation +except ImportError: + print("Error: This script requires PyObjC to be installed.") + print("Please install it with: pip install pyobjc") + sys.exit(1) + +# Constants for accessibility API +kAXErrorSuccess = 0 +kAXRoleAttribute = "AXRole" +kAXSubroleAttribute = "AXSubrole" +kAXTitleAttribute = "AXTitle" +kAXPositionAttribute = "AXPosition" +kAXSizeAttribute = "AXSize" +kAXChildrenAttribute = "AXChildren" +kAXMenuBarAttribute = "AXMenuBar" + + +def element_attribute(element, attribute): + """Get an attribute from an accessibility element""" + if attribute == kAXChildrenAttribute: + err, value = AXUIElementCopyAttributeValues(element, attribute, 0, 999, None) + if err == kAXErrorSuccess: + if isinstance(value, Foundation.NSArray): + return list(value) + else: + return value + err, value = AXUIElementCopyAttributeValue(element, attribute, None) + if err == kAXErrorSuccess: + return value + return None + + +def element_value(element, type): + """Get a value from an accessibility element""" + err, value = AXValueGetValue(element, type, None) + if err == True: + return value + return None + + +def get_element_bounds(element): + """Get the bounds of an accessibility element""" + bounds = { + "x": 0, + "y": 0, + "width": 0, + "height": 0 + } + + # Get position + position_value = element_attribute(element, kAXPositionAttribute) + if position_value: + position_value = element_value(position_value, kAXValueCGPointType) + if position_value: + bounds["x"] = position_value.x + bounds["y"] = position_value.y + + # Get size + size_value = element_attribute(element, kAXSizeAttribute) + if size_value: + size_value = element_value(size_value, kAXValueCGSizeType) + if size_value: + bounds["width"] = size_value.width + bounds["height"] = size_value.height + + return bounds + + +def find_dock_process(): + """Find the Dock process""" + running_apps = NSWorkspace.sharedWorkspace().runningApplications() + for app in running_apps: + if app.localizedName() == "Dock" and app.bundleIdentifier() == "com.apple.dock": + return app.processIdentifier() + return None + + +def get_menubar_bounds(): + """Get the bounds of the macOS menubar + + Returns: + Dictionary with x, y, width, height of the menubar + """ + # Get the system-wide accessibility element + system_element = AXUIElementCreateSystemWide() + + # Try to find the menubar + menubar = element_attribute(system_element, kAXMenuBarAttribute) + if menubar is None: + # If we can't get it directly, try through the frontmost app + frontmost_app = NSWorkspace.sharedWorkspace().frontmostApplication() + if frontmost_app: + app_pid = frontmost_app.processIdentifier() + app_element = AXUIElementCreateApplication(app_pid) + menubar = element_attribute(app_element, kAXMenuBarAttribute) + + if menubar is None: + print("Error: Could not get menubar") + # Return default menubar bounds as fallback + return {"x": 0, "y": 0, "width": 1800, "height": 24} + + # Get menubar bounds + return get_element_bounds(menubar) + + +def get_dock_bounds(): + """Get the bounds of the macOS Dock + + Returns: + Dictionary with x, y, width, height of the Dock + """ + dock_pid = find_dock_process() + if dock_pid is None: + print("Error: Could not find Dock process") + # Return empty bounds as fallback + return {"x": 0, "y": 0, "width": 0, "height": 0} + + # Create an accessibility element for the Dock + dock_element = AXUIElementCreateApplication(dock_pid) + if dock_element is None: + print(f"Error: Could not create accessibility element for Dock (PID {dock_pid})") + return {"x": 0, "y": 0, "width": 0, "height": 0} + + # Get the Dock's children + children = element_attribute(dock_element, kAXChildrenAttribute) + if not children or len(children) == 0: + print("Error: Could not get Dock children") + return {"x": 0, "y": 0, "width": 0, "height": 0} + + # Find the Dock's list (first child is usually the main dock list) + dock_list = None + for child in children: + role = element_attribute(child, kAXRoleAttribute) + if role == "AXList": + dock_list = child + break + + if dock_list is None: + print("Error: Could not find Dock list") + return {"x": 0, "y": 0, "width": 0, "height": 0} + + # Get the bounds of the dock list + return get_element_bounds(dock_list) + + +def get_ui_element_bounds(): + """Get the bounds of important UI elements like menubar and dock + + Returns: + Dictionary with menubar and dock bounds + """ + menubar_bounds = get_menubar_bounds() + dock_bounds = get_dock_bounds() + + return { + "menubar": menubar_bounds, + "dock": dock_bounds + } + + +if __name__ == "__main__": + # Example usage + bounds = get_ui_element_bounds() + print("Menubar bounds:", bounds["menubar"]) + print("Dock bounds:", bounds["dock"])