[feat]: add imperative environment customization support (#231)

Author: n8chur

Diff for: ViewEnvironmentUI/Sources/ViewEnvironmentPropagating.swift

@@ -119,6 +119,10 @@ extension ViewEnvironmentPropagating {
     public var environment: ViewEnvironment {
         var environment = environmentAncestor?.environment ?? .empty
 
+        for storedCustomization in customizations {
+            storedCustomization.customization(&environment)
+        }
+
         if let observing = self as? ViewEnvironmentObserving {
             observing.customize(environment: &environment)
         }
@@ -167,7 +171,7 @@ extension ViewEnvironmentPropagating {
     ///
     @_spi(ViewEnvironmentWiring)
     public func addEnvironmentNeedsUpdateObserver(
-        _ onNeedsUpdate: @escaping (ViewEnvironment) -> Void
+        _ onNeedsUpdate: @escaping ViewEnvironmentUpdateObservation
     ) -> ViewEnvironmentUpdateObservationLifetime {
         let object = ViewEnvironmentUpdateObservationKey()
         needsUpdateObservers[object] = onNeedsUpdate
@@ -176,6 +180,31 @@ extension ViewEnvironmentPropagating {
         }
     }
 
+    /// Adds a `ViewEnvironment` customization to this node.
+    ///
+    /// These customizations will occur before the node's `customize(environment:)` in cases where
+    /// this node conforms to `ViewEnvironmentObserving`, and will occur the order in which they
+    /// were added.
+    ///
+    /// The customization will only be active for as long as the returned lifetime is retained or
+    /// until `remove()` is called on it.
+    ///
+    @_spi(ViewEnvironmentWiring)
+    public func addEnvironmentCustomization(
+        _ customization: @escaping ViewEnvironmentCustomization
+    ) -> ViewEnvironmentCustomizationLifetime {
+        let storedCustomization = StoredViewEnvironmentCustomization(customization: customization)
+        customizations.append(storedCustomization)
+        return .init { [weak self] in
+            guard let self,
+                let index = self.customizations.firstIndex(where: { $0 === storedCustomization })
+            else {
+                return
+            }
+            self.customizations.remove(at: index)
+        }
+    }
+
     /// The `ViewEnvironment` propagation ancestor.
     ///
     /// This describes the ancestor that the `ViewEnvironment` is inherited from.
@@ -348,20 +377,26 @@ public typealias ViewEnvironmentUpdateObservation = (ViewEnvironment) -> Void
 public final class ViewEnvironmentUpdateObservationLifetime {
     /// Removes the observation.
     ///
-    /// This is called in `deinit`.
+    /// The observation is removed when the lifetime is de-initialized if this function was not
+    /// called before then.
     ///
     public func remove() {
+        guard let onRemove else {
+            assertionFailure("Environment update observation was already removed")
+            return
+        }
+        self.onRemove = nil
         onRemove()
     }
 
-    private let onRemove: () -> Void
+    private var onRemove: (() -> Void)?
 
     init(onRemove: @escaping () -> Void) {
         self.onRemove = onRemove
     }
 
     deinit {
-        remove()
+        onRemove?()
     }
 }
 
@@ -370,6 +405,7 @@ private enum ViewEnvironmentPropagatingNSObjectAssociatedKeys {
     static var needsUpdateObservers = NSObject()
     static var ancestorOverride = NSObject()
     static var descendantsOverride = NSObject()
+    static var customizations = NSObject()
 }
 
 extension ViewEnvironmentPropagating {
@@ -432,3 +468,68 @@ extension ViewEnvironmentPropagating {
 }
 
 private class ViewEnvironmentUpdateObservationKey: NSObject {}
+
+/// A closure that customizes the `ViewEnvironment` as it flows through a propagation node.
+///
+public typealias ViewEnvironmentCustomization = (inout ViewEnvironment) -> Void
+
+/// Describes the lifetime of a `ViewEnvironment` customization.
+///
+/// This customization will be removed when `remove()` is called or the lifetime token is
+/// de-initialized.
+///
+/// ## SeeAlso ##
+/// - `addEnvironmentCustomization(_:)`
+///
+public final class ViewEnvironmentCustomizationLifetime {
+    /// Removes the customization.
+    ///
+    /// The customization is removed when the lifetime is de-initialized if this function was not
+    /// called before then.
+    ///
+    public func remove() {
+        guard let onRemove else {
+            assertionFailure("Environment customization was already removed")
+            return
+        }
+        self.onRemove = nil
+        onRemove()
+    }
+
+    private var onRemove: (() -> Void)?
+
+    init(onRemove: @escaping () -> Void) {
+        self.onRemove = onRemove
+    }
+
+    deinit {
+        onRemove?()
+    }
+}
+
+extension ViewEnvironmentPropagating {
+    fileprivate var customizations: [StoredViewEnvironmentCustomization] {
+        get {
+            objc_getAssociatedObject(
+                self,
+                &AssociatedKeys.customizations
+            ) as? [StoredViewEnvironmentCustomization] ?? []
+        }
+        set {
+            objc_setAssociatedObject(
+                self,
+                &AssociatedKeys.customizations,
+                newValue,
+                .OBJC_ASSOCIATION_RETAIN_NONATOMIC
+            )
+        }
+    }
+}
+
+private final class StoredViewEnvironmentCustomization {
+    var customization: ViewEnvironmentCustomization
+
+    init(customization: @escaping ViewEnvironmentCustomization) {
+        self.customization = customization
+    }
+}

Diff for: ViewEnvironmentUI/Tests/ViewEnvironmentObservingTests.swift

@@ -191,7 +191,7 @@ final class ViewEnvironmentObservingTests: XCTestCase {
 
         var leafEnvironmentDidChangeCallCount = 0
         let leafNode = ViewEnvironmentPropagationNode(
-            environmentAncestor:  { [weak viewController] in
+            environmentAncestor: { [weak viewController] in
                 viewController
             },
             environmentDidChange: { _ in
@@ -226,22 +226,22 @@ final class ViewEnvironmentObservingTests: XCTestCase {
 
     func test_descendant_customFlow() {
         let descendant = TestViewEnvironmentObservingViewController()
-        
+
         let viewController = TestViewEnvironmentObservingViewController()
         viewController.environmentDescendantsOverride = { [descendant] }
-        
+
         viewController.applyEnvironmentIfNeeded()
         descendant.applyEnvironmentIfNeeded()
         XCTAssertFalse(viewController.needsEnvironmentUpdate)
         XCTAssertFalse(descendant.needsEnvironmentUpdate)
-        
+
         // With no ancestor configured the descendant should not respond to needing update
         viewController.setNeedsEnvironmentUpdate()
         XCTAssertTrue(viewController.needsEnvironmentUpdate)
         XCTAssertFalse(descendant.needsEnvironmentUpdate)
-        
+
         // With an ancestor defined the VC should respond to needing update
-        
+
         descendant.environmentAncestorOverride = { [weak viewController] in
             viewController
         }
@@ -318,12 +318,97 @@ final class ViewEnvironmentObservingTests: XCTestCase {
         XCTAssertEqual(observedEnvironments.count, 2)
         XCTAssertEqual(expectedTestContext, observedEnvironments.last?.testContext)
 
-        _ = observation // Suppress warning about variable never being read
+        withExtendedLifetime(observation) {}
         observation = nil
 
         container.setNeedsEnvironmentUpdate()
         XCTAssertEqual(observedEnvironments.count, 2)
     }
+
+    // MARK: - Customizations
+
+    func test_customization() throws {
+        let viewController = TestViewEnvironmentObservingViewController()
+
+        // Customizations should be respected as long as the lifetime exists.
+        do {
+            var customizationLifetime: ViewEnvironmentCustomizationLifetime? = viewController
+                .addEnvironmentCustomization {
+                    $0.testContext.number = 200
+                }
+
+            XCTAssertEqual(viewController.environment.testContext.number, 200)
+
+            withExtendedLifetime(customizationLifetime) {}
+            customizationLifetime = nil
+
+            // Customization should be removed when lifetime is deallocated
+            XCTAssertEqual(
+                viewController.environment.testContext.number,
+                TestContextKey.defaultValue.number
+            )
+        }
+
+        // Customizations should be respected until `remove()` is called on the lifetime.
+        do {
+            var customizationLifetime: ViewEnvironmentCustomizationLifetime? = viewController
+                .addEnvironmentCustomization {
+                    $0.testContext.number = 200
+                }
+
+            XCTAssertEqual(viewController.environment.testContext.number, 200)
+
+            customizationLifetime?.remove()
+
+            // Customization should be removed when lifetime is deallocated
+            XCTAssertEqual(
+                viewController.environment.testContext.number,
+                TestContextKey.defaultValue.number
+            )
+
+            withExtendedLifetime(customizationLifetime) {}
+            customizationLifetime = nil
+        }
+
+        // Customizations should occur in the order they are added
+        do {
+            var customization1Lifetime: ViewEnvironmentCustomizationLifetime? = viewController
+                .addEnvironmentCustomization {
+                    $0.testContext.number = 100
+                }
+
+            var customization2Lifetime: ViewEnvironmentCustomizationLifetime? = viewController
+                .addEnvironmentCustomization {
+                    $0.testContext.number = 200
+                }
+
+            XCTAssertEqual(viewController.environment.testContext.number, 200)
+
+            withExtendedLifetime(customization1Lifetime) {}
+            customization1Lifetime = nil
+            withExtendedLifetime(customization2Lifetime) {}
+            customization2Lifetime = nil
+        }
+
+        // Customizations should favor the nodes customizations if present
+        do {
+            viewController.customizeEnvironment = {
+                $0.testContext.number = 300
+            }
+
+            var customizationLifetime: ViewEnvironmentCustomizationLifetime? = viewController
+                .addEnvironmentCustomization {
+                    $0.testContext.number = 200
+                    $0.testContext.string = "200"
+                }
+
+            XCTAssertEqual(viewController.environment.testContext.number, 300)
+            XCTAssertEqual(viewController.environment.testContext.string, "200")
+
+            withExtendedLifetime(customizationLifetime) {}
+            customizationLifetime = nil
+        }
+    }
 }
 
 // MARK: - Helpers