UIViewController+Screenshot.swift

//
//  UIViewController+Screenshot.swift
//  AppleLibs
//
//  Created by Tobias on 29.04.21.
//

import UIKit

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