IdentityFlatten#when and IdentityFlatten#unless. (#1453)

zio · Jan 15, 2025 · a69f4c3 · a69f4c3
1 parent ce9ee2f
commit a69f4c3
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 0 deletions.
diff --git a/core-tests/shared/src/test/scala/zio/prelude/IdentityFlattenSpec.scala b/core-tests/shared/src/test/scala/zio/prelude/IdentityFlattenSpec.scala
@@ -7,6 +7,15 @@ import zio.test.laws._
 object IdentityFlattenSpec extends ZIOBaseSpec {
   import zio.prelude.Fixtures._
 
+  val genBoolean: Gen[Any, Boolean] =
+    Gen.boolean
+
+  val genInt: Gen[Any, Int] =
+    Gen.int
+
+  val genList: Gen[Sized, List[Int]] =
+    Gen.listOf(genInt)
+
   def spec: Spec[Environment, Any] =
     suite("IdentityFlattenSpec")(
       suite("laws")(
@@ -16,6 +25,22 @@ object IdentityFlattenSpec extends ZIOBaseSpec {
         test("option")(checkAllLaws(IdentityFlattenLaws)(GenF.option, Gen.int)),
         test("optional")(checkAllLaws(IdentityFlattenLaws)(optionalGenF, Gen.int)),
         test("vector")(checkAllLaws(IdentityFlattenLaws)(GenF.vector, Gen.int))
+      ),
+      suite("combinators")(
+        test("when") {
+          check(genList, genBoolean) { (as, b) =>
+            val actual   = as.when(b)
+            val expected = if (b) as.map(Some(_)) else List(None)
+            assert(actual)(equalTo(expected))
+          }
+        },
+        test("unless") {
+          check(genList, genBoolean) { (as, b) =>
+            val actual   = as.unless(b)
+            val expected = if (b) List(None) else as.map(Some(_))
+            assert(actual)(equalTo(expected))
+          }
+        }
       )
     )
 

diff --git a/core/shared/src/main/scala/zio/prelude/IdentityFlatten.scala b/core/shared/src/main/scala/zio/prelude/IdentityFlatten.scala
@@ -43,3 +43,17 @@ object IdentityFlatten {
     identityFlatten
 
 }
+
+trait IdentityFlattenSyntax {
+
+  /**
+   * Provides infix syntax for identity operations for covariant types.
+   */
+  implicit class IdentityFlattenCovariantOps[F[+_], A](fa: F[A]) {
+    def unless(b: => Boolean)(implicit identity: IdentityFlatten[F], covariant: Covariant[F]): F[Option[A]] =
+      if (b) identity.any.as(None) else fa.map(Some(_))
+
+    def when(b: => Boolean)(implicit identity: IdentityFlatten[F], covariant: Covariant[F]): F[Option[A]] =
+      if (b) fa.map(Some(_)) else identity.any.as(None)
+  }
+}
diff --git a/core/shared/src/main/scala/zio/prelude/package.scala b/core/shared/src/main/scala/zio/prelude/package.scala
@@ -38,6 +38,7 @@ package object prelude
     with IdExports
     with IdentityBothSyntax
     with IdentityEitherSyntax
+    with IdentityFlattenSyntax
     with IdentitySyntax
     with InvariantSyntax
     with InverseSyntax