From 587ce77f051556f5dc704d9174b5b96d580d8e63 Mon Sep 17 00:00:00 2001 From: Varun Nuthalapati Date: Tue, 28 Apr 2026 19:37:14 -0700 Subject: [PATCH] fix: guard against None control in ControlPhotographer.capture When the application window becomes unavailable (e.g. it closes between the screenshot request and the actual capture), the control element passed to ControlPhotographer can be None. Calling .capture_as_image() on None raises: AttributeError: 'NoneType' object has no attribute 'capture_as_image' Add an explicit None-check that logs a warning and returns None, letting callers handle the missing screenshot gracefully instead of crashing. Closes #232 --- ufo/automator/ui_control/screenshot.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ufo/automator/ui_control/screenshot.py b/ufo/automator/ui_control/screenshot.py index 34af17ee..12f25384 100644 --- a/ufo/automator/ui_control/screenshot.py +++ b/ufo/automator/ui_control/screenshot.py @@ -79,12 +79,17 @@ def __init__(self, control: UIAWrapper): """ self.control = control - def capture(self, save_path: str = None, scalar: List[int] = None) -> Image.Image: + def capture(self, save_path: str = None, scalar: List[int] = None) -> Optional[Image.Image]: """ Capture a screenshot. :param save_path: The path to save the screenshot. - :return: The screenshot. + :param scalar: Optional rescale dimensions [width, height]. + :return: The screenshot, or None if the control is unavailable. """ + if self.control is None: + logger.warning("Cannot capture screenshot: control element is None.") + return None + # Capture single window screenshot screenshot = self.control.capture_as_image() if scalar is not None: