expand produce costuzation/

trafi · Jan 2, 2024 · 04fdb6c · 04fdb6c
1 parent 9388a9e
commit 04fdb6c
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 8 deletions.
diff --git a/Sources/StoryFlow/Interfaces/Helpers/Testing.swift b/Sources/StoryFlow/Interfaces/Helpers/Testing.swift
@@ -11,12 +11,27 @@ extension InputRequiring where Self: UIViewController {
 
 extension OutputProducing where Self: UIViewController {
 
-    public init(nibName: String? = nil, bundle: Bundle? = nil, produce: @escaping (OutputType) -> ()) {
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        produce: @escaping (OutputType) -> OutputType?
+    ) {
         self.init(nibName: nibName, bundle: bundle)
         produceStub = produce
     }
+
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        produce: @escaping (OutputType) -> ()
+    ) {
+        self.init(nibName: nibName, bundle: bundle) {
+            produce($0)
+            return nil
+        }
+    }
 
-    var produceStub: ((OutputType) -> ())? {
+    var produceStub: ((OutputType) -> OutputType?)? {
         get { return associated(with: &produceStubKey) }
         set { associate(newValue, with: &produceStubKey) }
     }
@@ -25,10 +40,27 @@ private var produceStubKey = 0
 
 extension InputRequiring where Self: UIViewController & OutputProducing {
 
-    public init(nibName: String? = nil, bundle: Bundle? = nil,
-                input: InputType, produce: @escaping (OutputType) -> ()) {
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        input: InputType,
+        produce: @escaping (OutputType) -> OutputType?
+    ) {
         self.init(nibName: nibName, bundle: bundle)
         self.input = input
         self.produceStub = produce
     }
+
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        input: InputType,
+        produce: @escaping (OutputType) -> ()
+    ) {
+        self.init(nibName: nibName, bundle: bundle) {
+            produce($0)
+            return nil
+        }
+    }
+
 }
diff --git a/Sources/StoryFlow/Interfaces/OuputProducing.swift b/Sources/StoryFlow/Interfaces/OuputProducing.swift
@@ -19,10 +19,17 @@ extension OutputProducing where Self: UIViewController {
      - Parameter output: The value being passed to next view controller.
      */
     public func produce(_ output: OutputType) {
-
-        if let produce = produceStub {
-            produce(output)
-        } else if let flow = flow {
+        if let produceStub {
+            if let customOutput = produceStub(output) {
+                _produce(customOutput)
+            }
+        } else {
+            _produce(output)
+        }
+    }
+
+    private func _produce(_ output: OutputType) {
+        if let flow = flow {
             flow.proceed(with: output, from: self)
         } else {
             implicitFlow.proceed(with: output, from: self)