Internal

PiperOrigin-RevId: 476134510
toptal · Nov 8, 2022 · 88905ef · 88905ef
1 parent 3fed2a8
commit 88905ef
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 28 deletions.
diff --git a/UILib/GREYScreenshotter.m b/UILib/GREYScreenshotter.m
@@ -33,18 +33,6 @@
 // Private class for ModalHostingWindow window that doesn't work with drawViewHierarchyInRect.
 static Class gUIModalItemHostingWindowClass;
 
-/**
- * @return A current screen if the window exists and @c nil if it does not.
- */
-static UIScreen *MainScreen(void) {
-  UIScreen *mainScreen = [GREYUILibUtils screen];
-  if (!mainScreen || CGRectEqualToRect(mainScreen.bounds, CGRectNull)) {
-    return nil;
-  } else {
-    return mainScreen;
-  }
-}
-
 @implementation GREYScreenshotter
 
 + (void)initialize {
@@ -60,8 +48,7 @@ + (void)drawScreenInContext:(CGContextRef)bitmapContextRef
               withStatusBar:(BOOL)includeStatusBar {
   GREYFatalAssertWithMessage(CGBitmapContextGetBitmapInfo(bitmapContextRef) != 0,
                              @"The context ref must point to a CGBitmapContext.");
-  UIScreen *mainScreen = MainScreen();
-  if (!mainScreen) return;
+  UIScreen *mainScreen = [GREYUILibUtils screen];
   CGRect screenRect = mainScreen.bounds;
   [self drawScreenInContext:bitmapContextRef
          afterScreenUpdates:afterUpdates
@@ -89,9 +76,7 @@ + (UIImage *)snapshotElement:(id)element {
                                ? (UIView *)element
                                : [element grey_viewContainingSelf];
 
-  UIScreen *mainScreen = MainScreen();
-  if (!mainScreen) return nil;
-  UIGraphicsBeginImageContextWithOptions(elementAXFrame.size, NO, mainScreen.scale);
+  UIGraphicsBeginImageContextWithOptions(elementAXFrame.size, NO, [GREYUILibUtils screen].scale);
   [self drawViewInContext:UIGraphicsGetCurrentContext()
                      view:viewToSnapshot
                    bounds:elementAXFrame
@@ -111,9 +96,7 @@ + (NSString *)saveImageAsPNG:(UIImage *)image
 
 + (UIImage *)grey_takeScreenshotAfterScreenUpdates:(BOOL)afterScreenUpdates
                                      withStatusBar:(BOOL)includeStatusBar {
-  UIScreen *mainScreen = MainScreen();
-  if (!mainScreen) return nil;
-  CGRect screenRect = mainScreen.bounds;
+  CGRect screenRect = [GREYUILibUtils screen].bounds;
   return [self grey_takeScreenshotAfterScreenUpdates:afterScreenUpdates
                                         inScreenRect:screenRect
                                        withStatusBar:includeStatusBar];
@@ -169,8 +152,8 @@ + (void)drawViewInContext:(CGContextRef)bitmapContextRef
                      view:(UIView *)view
                    bounds:(CGRect)boundsInScreenRect
        afterScreenUpdates:(BOOL)afterScreenUpdates {
-  UIScreen *mainScreen = MainScreen();
-  if (!mainScreen) return;
+  UIScreen *mainScreen = [GREYUILibUtils screen];
+
   // The bitmap context width and height are scaled, so we need to undo the scale adjustment.
   CGFloat scale = mainScreen.scale;
   CGFloat contextWidth = CGBitmapContextGetWidth(bitmapContextRef) / scale;

diff --git a/UILib/Utils/GREYUILibUtils.m b/UILib/Utils/GREYUILibUtils.m
@@ -15,7 +15,6 @@
 //
 
 #import "GREYUILibUtils.h"
-#import "GREYLogger.h"
 
 UIWindow *GREYUILibUtilsGetApplicationKeyWindow(UIApplication *application) {
   // New API only available on Xcode 13+
@@ -71,11 +70,6 @@ + (UIScreen *)screen {
   UIScreen *screen;
 
   if (@available(iOS 13.0, *)) {
-    UIWindow *window = [self window];
-    if (!window) {
-      GREYLog(@"Attempted to take a screenshot when no window exists, i.e. window == nil");
-      return nil;
-    }
     screen = [self window].windowScene.screen;
     // This check is added in case there is an issue with getting the screen i.e. if the screen
     // hasn't come up.