Split functions (#102)

EmergeTools · Feb 28, 2024 · dfde4c5 · dfde4c5
1 parent 1ab7f66
commit dfde4c5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 11 deletions.
diff --git a/Sources/SnapshotPreviewsCore/View+Snapshot.swift b/Sources/SnapshotPreviewsCore/View+Snapshot.swift
@@ -16,23 +16,28 @@ public enum RenderingError: Error {
 }
 
 extension View {
-  public func snapshot(
-    layout: PreviewLayout,
-    window: UIWindow,
-    async: Bool,
-    completion: @escaping (Result<UIImage, RenderingError>, Float?, Bool?) -> Void)
-  {
+  public func makeExpandingView(layout: PreviewLayout, window: UIWindow) -> ExpandingViewController {
     UIView.setAnimationsEnabled(false)
     let animationDisabledView = self.transaction { transaction in
       transaction.disablesAnimations = true
     }
     let controller = ExpandingViewController(rootView: animationDisabledView)
     controller.setupView(layout: layout)
 
-    let (windowRootVC, containerVC) = Self.setupRootVC(subVC: controller)
+    let windowRootVC = Self.setupRootVC(subVC: controller)
     window.rootViewController = windowRootVC
-    controller.expansionSettled = { [weak containerVC, weak controller] renderingMode, precision, accessibilityEnabled in
-      guard let containerVC, let controller else { return }
+    return controller
+  }
+
+  public func snapshot(
+    layout: PreviewLayout,
+    controller: ExpandingViewController,
+    window: UIWindow,
+    async: Bool,
+    completion: @escaping (Result<UIImage, RenderingError>, Float?, Bool?) -> Void)
+  {
+    controller.expansionSettled = { [weak controller, weak window] renderingMode, precision, accessibilityEnabled in
+      guard let controller, let window, let containerVC = controller.parent else { return }
 
       if async {
         DispatchQueue.main.asyncAfter(deadline: .now() + 2) {
@@ -71,7 +76,7 @@ extension View {
     }
   }
 
-  private static func setupRootVC(subVC: UIViewController) -> (UIViewController, UIViewController) {
+  private static func setupRootVC(subVC: UIViewController) -> UIViewController {
     let windowRootVC = UIViewController()
     windowRootVC.view.bounds = UIScreen.main.bounds
     windowRootVC.view.backgroundColor = .clear
@@ -96,7 +101,7 @@ extension View {
     subVC.view.widthAnchor.constraint(lessThanOrEqualToConstant: UIScreen.main.bounds.width).isActive = true
     containerVC.view.heightAnchor.constraint(greaterThanOrEqualTo: subVC.view.heightAnchor, multiplier: 1).isActive = true
 
-    return (windowRootVC, containerVC)
+    return windowRootVC
   }
 
   private static func takeSnapshot(

diff --git a/Sources/SnapshottingSwift/Snapshots.swift b/Sources/SnapshottingSwift/Snapshots.swift
@@ -125,8 +125,10 @@ class Snapshots {
   @MainActor func display(preview: SnapshotPreviewsCore.Preview, completion: @escaping (Result<UIImage, RenderingError>, Float?, Bool?) -> Void) throws {
     var view = preview.view()
     view = PreferredColorSchemeWrapper { AnyView(view) }
+    let controller = view.makeExpandingView(layout: preview.layout, window: window)
     view.snapshot(
       layout: preview.layout,
+      controller: controller,
       window: window,
       async: false,
       completion: completion)