Add Output.when helper method

fixes #435
VirtusLab · Apr 5, 2024 · 2d062b9 · 2d062b9
1 parent ad94f26
commit 2d062b9
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 0 deletions.
diff --git a/core/src/main/scala/besom/internal/Output.scala b/core/src/main/scala/besom/internal/Output.scala
@@ -80,6 +80,11 @@ trait OutputFactory:
   )(using BuildFrom[CC[Output[B]], B, To], Context): Output[To] = sequence(coll.map(f).asInstanceOf[CC[Output[B]]])
 
   def fail(t: Throwable)(using Context): Output[Nothing] = Output.fail(t)
+
+  def when[A: Typeable](cond: => Boolean | Output[Boolean] | Output[Option[Boolean]])(
+    a: => A | Output[A] | Output[Option[A]]
+  )(using ctx: Context): Output[Option[A]] = Output.when(cond)(a)
+
 trait OutputExtensionsFactory:
   implicit final class OutputSequenceOps[A, CC[X] <: Iterable[X], To](coll: CC[Output[A]]):
     def sequence(using BuildFrom[CC[Output[A]], A, To], Context): Output[To] =
@@ -159,4 +164,31 @@ object Output:
   def secret[A](value: A)(using ctx: Context): Output[A] =
     new Output[A](ctx.registerTask(Result.pure(OutputData(value, Set.empty, isSecret = true))))
 
+  def when[A: Typeable](cond: => Boolean | Option[Boolean] | Output[Boolean] | Output[Option[Boolean]])(
+    a: => A | Option[A] | Output[A] | Output[Option[A]]
+  )(using ctx: Context): Output[Option[A]] =
+    val p: Output[Boolean] = cond match
+      case o: Output[Boolean | Option[Boolean]] =>
+        o.flatMap {
+          case None             => Output(false)
+          case Some(b: Boolean) => Output(b)
+          case b: Boolean       => Output(b)
+        }
+      case b: Option[Boolean] => Output(b.getOrElse(false))
+      case b: Boolean         => Output(b)
+
+    def f(c: Boolean): Output[Option[A]] =
+      a match
+        case o: Output[A @unchecked] if c =>
+          o.flatMap {
+            case None       => Output(None)
+            case Some(v: A) => Output(Some(v))
+            case a: A       => Output(Some(a))
+          }
+        case v: Option[A @unchecked] if c => Output(v)
+        case v: A if c                    => Output(Some(v))
+        case _                            => Output(None)
+
+    p.flatMap(f)
+  end when
 end Output
diff --git a/core/src/test/scala/besom/internal/OutputTest.scala b/core/src/test/scala/besom/internal/OutputTest.scala
@@ -138,4 +138,46 @@ class OutputTest extends munit.FunSuite:
     Context().waitForAllTasks.unsafeRunSync()
   }
 
+  Vector(
+    (true, "value", Some("value")),
+    (false, "value", None)
+  ).foreach { (cond, value, expected) =>
+    given Context = DummyContext().unsafeRunSync()
+    for
+      optCond <- Vector(true, false)
+      outCond <- Vector(true, false)
+      optVal  <- Vector(true, false)
+      outVal  <- Vector(true, false)
+    do
+      val c: Boolean | Option[Boolean] | Output[Boolean] | Output[Option[Boolean]] = (outCond, optCond) match
+        case (true, true)   => Output(Option(cond))
+        case (true, false)  => Output(cond)
+        case (false, true)  => Some(cond)
+        case (false, false) => cond
+
+      val v: String | Option[String] | Output[String] | Output[Option[String]] = (outVal, optVal) match
+        case (true, true)   => Output(Option(value))
+        case (true, false)  => Output(value)
+        case (false, true)  => Some(value)
+        case (false, false) => value
+
+      test(s"when ${cond} then ${value} (optCond: ${optCond}, outCond: ${outCond}, valOpt: ${optVal}, valOut: ${outVal})") {
+        val result = Output.when(c)(v)
+        assertEquals(result.getData.unsafeRunSync(), OutputData(expected))
+      }
+
+    Context().waitForAllTasks.unsafeRunSync()
+  }
+
+  /*test("when short circuiting false condition") {
+    given Context = DummyContext().unsafeRunSync()
+
+    val v: Output[String] = Output({
+      fail("This should not be evaluated")
+      "value"
+    })
+    val result = Output.when(false)(v)
+    assertEquals(result.getData.unsafeRunSync(), OutputData(None))
+  }*/
+
 end OutputTest