Update screenshot function

AppleLibs/View/UIViewController+Screenshot.swift

@@ -10,24 +10,20 @@ import UIKit
 extension UIViewController
 {
     open func takeScreenshot(_ shouldSave: Bool = false) -> UIImage? {
-        var screenshotImage :UIImage?
-        
-        let scale = UIScreen.main.scale
-        guard let layer = self.view.window?.layer else {
+        guard let window = self.view.window else {
             return nil
         }
+        let bounds = UIScreen.main.bounds
         
-        UIGraphicsBeginImageContextWithOptions(layer.frame.size, false, scale);
-        guard let context = UIGraphicsGetCurrentContext() else {
+        UIGraphicsBeginImageContextWithOptions(bounds.size, false, 0)
+        window.drawHierarchy(in: bounds, afterScreenUpdates: true)
+        guard let screenshotImage = UIGraphicsGetImageFromCurrentImageContext() else {
             return nil
         }
-        
-        layer.render(in:context)
-        screenshotImage = UIGraphicsGetImageFromCurrentImageContext()
         UIGraphicsEndImageContext()
         
-        if let image = screenshotImage, shouldSave {
-            UIImageWriteToSavedPhotosAlbum(image, nil, nil, nil)
+        if shouldSave {
+            UIImageWriteToSavedPhotosAlbum(screenshotImage, nil, nil, nil)
         }
         return screenshotImage
     }