zio · iravid · May 23, 2020 · May 20, 2020 · May 20, 2020 · May 20, 2020
diff --git a/streams-tests/shared/src/test/scala/zio/stream/ZStreamSpec.scala b/streams-tests/shared/src/test/scala/zio/stream/ZStreamSpec.scala
@@ -17,7 +17,6 @@ import zio.test.environment.TestClock
 
 object ZStreamSpec extends ZIOBaseSpec {
   import ZIOTag._
-  import ZStream.Take
 
   def inParallel(action: => Unit)(implicit ec: ExecutionContext): Unit =
     ec.execute(() => action)
@@ -755,7 +754,7 @@ object ZStreamSpec extends ZIOBaseSpec {
               ZStream
                 .range(0, 10)
                 .toQueue(1)
-                .use(q => ZStream.fromQueue(q).collectWhileSuccess.runCollect)
+                .use(q => ZStream.fromQueue(q).map(_.exit).collectWhileSuccess.runCollect)
                 .map(_.flatMap(_.toList))
             )(equalTo(Range(0, 10).toList))
           },
@@ -764,7 +763,7 @@ object ZStreamSpec extends ZIOBaseSpec {
             assertM(
               (ZStream.range(0, 10) ++ ZStream.fail(e))
                 .toQueue(1)
-                .use(q => ZStream.fromQueue(q).collectWhileSuccess.runCollect)
+                .use(q => ZStream.fromQueue(q).map(_.exit).collectWhileSuccess.runCollect)
                 .run
             )(fails(equalTo(e)))
           } @@ zioTag(errors)
@@ -1298,25 +1297,25 @@ object ZStreamSpec extends ZIOBaseSpec {
             assertM(
               ZStream
                 .fromChunks(chunks: _*)
-                .mapChunks(chunk => Chunk.single(Exit.succeed(chunk)))
+                .mapChunks(chunk => Chunk.single(Take.chunk(chunk)))
                 .flattenTake
                 .runCollect
             )(equalTo(chunks.fold(Chunk.empty)(_ ++ _).toList))
           }),
           testM("stop collecting on Exit.Failure") {
             assertM(
               ZStream(
-                Exit.succeed(Chunk(1, 2)),
-                Exit.succeed(Chunk.single(3)),
-                Exit.fail(Option.empty[Unit])
+                Take.chunk(Chunk(1, 2)),
+                Take.single(3),
+                Take.end
               ).flattenTake.runCollect
             )(equalTo(List(1, 2, 3)))
           },
           testM("work with empty chunks") {
-            assertM(ZStream(Exit.succeed(Chunk.empty), Exit.succeed(Chunk.empty)).flattenTake.runCollect)(isEmpty)
+            assertM(ZStream(Take.chunk(Chunk.empty), Take.chunk(Chunk.empty)).flattenTake.runCollect)(isEmpty)
           },
           testM("work with empty streams") {
-            assertM(ZStream.fromIterable[Exit[Option[Nothing], Chunk[Nothing]]](Nil).flattenTake.runCollect)(isEmpty)
+            assertM(ZStream.fromIterable[Take[Nothing, Nothing]](Nil).flattenTake.runCollect)(isEmpty)
           }
         ),
         suite("foreach")(
@@ -2546,15 +2545,15 @@ object ZStreamSpec extends ZIOBaseSpec {
               s.toQueue(1000)
                 .use(queue => queue.size.repeat(Schedule.doWhile(_ != c.size + 1)) *> queue.takeAll)
             )(
-              equalTo(c.toSeq.toList.map(i => Exit.succeed(Chunk(i))) :+ Exit.fail(None))
+              equalTo(c.toSeq.toList.map(Take.single) :+ Take.end)
             )
           }),
           testM("toQueueUnbounded")(checkM(Gen.chunkOfBounded(0, 3)(Gen.anyInt)) { (c: Chunk[Int]) =>
             val s = ZStream.fromChunk(c).flatMap(ZStream.succeed(_))
             assertM(
               s.toQueueUnbounded.use(queue => queue.size.repeat(Schedule.doWhile(_ != c.size + 1)) *> queue.takeAll)
             )(
-              equalTo(c.toSeq.toList.map(i => Exit.succeed(Chunk(i))) :+ Exit.fail(None))
+              equalTo(c.toSeq.toList.map(Take.single) :+ Take.end)
             )
           })
         ),
@@ -2853,17 +2852,16 @@ object ZStreamSpec extends ZIOBaseSpec {
           }
         },
         testM("fromTQueue") {
-          TQueue.bounded[Int](5).commit.flatMap {
-            tqueue =>
-              ZStream.fromTQueue(tqueue).toQueueUnbounded.use { queue =>
-                for {
-                  _      <- tqueue.offerAll(List(1, 2, 3)).commit
-                  first  <- ZStream.fromQueue(queue).take(3).runCollect
-                  _      <- tqueue.offerAll(List(4, 5)).commit
-                  second <- ZStream.fromQueue(queue).take(2).runCollect
-                } yield assert(first)(equalTo(List(1, 2, 3).map(i => Exit.succeed(Chunk.single(i))))) &&
-                  assert(second)(equalTo(List(4, 5).map(i => Exit.succeed(Chunk.single(i)))))
-              }
+          TQueue.bounded[Int](5).commit.flatMap { tqueue =>
+            ZStream.fromTQueue(tqueue).toQueueUnbounded.use { queue =>
+              for {
+                _      <- tqueue.offerAll(List(1, 2, 3)).commit
+                first  <- ZStream.fromQueue(queue).take(3).runCollect
+                _      <- tqueue.offerAll(List(4, 5)).commit
+                second <- ZStream.fromQueue(queue).take(2).runCollect
+              } yield assert(first)(equalTo(List(1, 2, 3).map(Take.single))) &&
+                assert(second)(equalTo(List(4, 5).map(Take.single)))
+            }
           }
         } @@ flaky,
         testM("iterate")(
@@ -2963,7 +2961,7 @@ object ZStreamSpec extends ZIOBaseSpec {
     ) @@ TestAspect.timed
 
   trait ChunkCoordination[A] {
-    def queue: Queue[Take[Nothing, A]]
+    def queue: Queue[Exit[Option[Nothing], Chunk[A]]]
     def offer: UIO[Boolean]
     def proceed: UIO[Unit]
     def awaitNext: UIO[Unit]
@@ -2973,13 +2971,13 @@ object ZStreamSpec extends ZIOBaseSpec {
     chunks: List[Chunk[A]]
   )(assertion: ChunkCoordination[A] => ZIO[Clock with TestClock, Nothing, TestResult]) =
     for {
-      q  <- Queue.unbounded[Take[Nothing, A]]
+      q  <- Queue.unbounded[Exit[Option[Nothing], Chunk[A]]]
       ps <- Queue.unbounded[Unit]
       ref <- Ref
-              .make[List[List[Take[Nothing, A]]]](
+              .make[List[List[Exit[Option[Nothing], Chunk[A]]]]](
                 chunks.init.map { chunk =>
                   List(Exit.succeed(chunk))
-                } ++ chunks.lastOption.map(chunk => List(Exit.succeed(chunk), Take.End))
+                } ++ chunks.lastOption.map(chunk => List(Exit.succeed(chunk), Exit.fail(None)))
               )
       chunkCoordination = new ChunkCoordination[A] {
         val queue = q

diff --git a/streams/jvm/src/main/scala/zio/stream/platform.scala b/streams/jvm/src/main/scala/zio/stream/platform.scala
@@ -56,12 +56,12 @@ trait ZStreamPlatformSpecificConstructors { self: ZStream.type =>
   ): ZStream[R, E, A] =
     ZStream {
       for {
-        output  <- Queue.bounded[Take[E, A]](outputBuffer).toManaged(_.shutdown)
+        output  <- Queue.bounded[stream.Take[E, A]](outputBuffer).toManaged(_.shutdown)
         runtime <- ZIO.runtime[R].toManaged_
         eitherStream <- ZManaged.effectTotal {
                          register(k =>
                            try {
-                             runtime.unsafeRun(k.run.flatMap(output.offer))
+                             runtime.unsafeRun(stream.Take.fromPull(k).flatMap(output.offer))
                              ()
                            } catch {
                              case FiberFailure(c) if c.interrupted =>
@@ -75,7 +75,7 @@ trait ZStreamPlatformSpecificConstructors { self: ZStream.type =>
                    } yield done.get.flatMap {
                      if (_) Pull.end
                      else
-                       output.take.flatMap(Pull.fromTake).onError(_ => done.set(true) *> output.shutdown)
+                       output.take.flatMap(_.done).onError(_ => done.set(true) *> output.shutdown)
                    }).ensuring(canceler)
                  case Right(stream) => output.shutdown.toManaged_ *> stream.process
                }
@@ -93,11 +93,11 @@ trait ZStreamPlatformSpecificConstructors { self: ZStream.type =>
   ): ZStream[R, E, A] =
     managed {
       for {
-        output  <- Queue.bounded[Take[E, A]](outputBuffer).toManaged(_.shutdown)
+        output  <- Queue.bounded[stream.Take[E, A]](outputBuffer).toManaged(_.shutdown)
         runtime <- ZIO.runtime[R].toManaged_
         _ <- register { k =>
               try {
-                runtime.unsafeRun(k.run.flatMap(output.offer))
+                runtime.unsafeRun(stream.Take.fromPull(k).flatMap(output.offer))
                 ()
               } catch {
                 case FiberFailure(c) if c.interrupted =>
@@ -108,7 +108,7 @@ trait ZStreamPlatformSpecificConstructors { self: ZStream.type =>
           if (_)
             Pull.end
           else
-            output.take.flatMap(Pull.fromTake).onError(_ => done.set(true) *> output.shutdown)
+            output.take.flatMap(_.done).onError(_ => done.set(true) *> output.shutdown)
         }
       } yield pull
     }.flatMap(repeatEffectChunkOption(_))
@@ -125,12 +125,12 @@ trait ZStreamPlatformSpecificConstructors { self: ZStream.type =>
   ): ZStream[R, E, A] =
     ZStream {
       for {
-        output  <- Queue.bounded[Take[E, A]](outputBuffer).toManaged(_.shutdown)
+        output  <- Queue.bounded[stream.Take[E, A]](outputBuffer).toManaged(_.shutdown)
         runtime <- ZIO.runtime[R].toManaged_
         maybeStream <- ZManaged.effectTotal {
                         register { k =>
                           try {
-                            runtime.unsafeRun(k.run.flatMap(output.offer))
+                            runtime.unsafeRun(stream.Take.fromPull(k).flatMap(output.offer))
                             ()
                           } catch {
                             case FiberFailure(c) if c.interrupted =>
@@ -146,7 +146,7 @@ trait ZStreamPlatformSpecificConstructors { self: ZStream.type =>
                      if (_)
                        Pull.end
                      else
-                       output.take.flatMap(Pull.fromTake).onError(_ => done.set(true) *> output.shutdown)
+                       output.take.flatMap(_.done).onError(_ => done.set(true) *> output.shutdown)
                    }
                }
       } yield pull

diff --git a/streams/shared/src/main/scala/zio/stream/Take.scala b/streams/shared/src/main/scala/zio/stream/Take.scala
@@ -0,0 +1,136 @@
+/*
+ * Copyright 2017-2020 John A. De Goes and the ZIO Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package zio.stream
+
+import zio._
+
+/**
+ * A `Take[E, A]` represents a single `take` from a queue modeling a stream of
+ * values. A `Take` may be a failure cause `Cause[E]`, an chunk value `A`
+ * or an end-of-stream marker.
+ */
+case class Take[+E, +A](exit: Exit[Option[E], Chunk[A]]) extends AnyVal {
+
+  /**
+   * Transforms `Take[E, A]` to `ZIO[R, E, B]`.
+   */
+  def done[R]: ZIO[R, Option[E], Chunk[A]] =
+    IO.done(exit)
+
+  /**
+   * Folds over the failure cause, success value and end-of-stream marker to
+   * yield a value.
+   */
+  def fold[Z](end: => Z, error: Cause[E] => Z, value: Chunk[A] => Z): Z =
+    exit.fold(Cause.sequenceCauseOption(_).fold(end)(error), value)
+
+  /**
+   * Effectful version of [[Take#fold]].
+   *
+   * Folds over the failure cause, success value and end-of-stream marker to
+   * yield an effect.
+   */
+  def foldM[R, E1, Z](
+    end: => ZIO[R, E1, Z],
+    error: Cause[E] => ZIO[R, E1, Z],
+    value: Chunk[A] => ZIO[R, E1, Z]
+  ): ZIO[R, E1, Z] =
+    exit.foldM(Cause.sequenceCauseOption(_).fold(end)(error), value)
+
+  /**
+   * Checks if this `take` is done (`Take.end`).
+   */
+  def isDone: Boolean =
+    exit.fold(Cause.sequenceCauseOption(_).isEmpty, _ => false)
+
+  /**
+   * Transforms `Take[E, A]` to `Take[E, B]` by applying function `f`.
+   */
+  def map[B](f: A => B): Take[E, B] =
+    Take(exit.map(_.map(f)))
+
+  /**
+   * Returns an effect that effectfully "peeks" at the success of this take.
+   */
+  def tap[R, E1](f: Chunk[A] => ZIO[R, E1, Any]): ZIO[R, E1, Unit] =
+    exit.foreach(f).unit
+}
+
+object Take {
+
+  /**
+   * Creates a `Take[Nothing, A]` with a singleton chunk.
+   */
+  def single[A](a: A): Take[Nothing, A] =
+    Take(Exit.succeed(Chunk.single(a)))
+
+  /**
+   * Creates a `Take[Nothing, A]` with the specified chunk.
+   */
+  def chunk[A](as: Chunk[A]): Take[Nothing, A] =
+    Take(Exit.succeed(as))
+
+  /**
+   * Creates a failing `Take[E, Nothing]` with the specified failure.
+   */
+  def fail[E](e: E): Take[E, Nothing] =
+    Take(Exit.fail(Some(e)))
+
+  /**
+   * Creates an effect from `ZIO[R, E,A]` that does not fail, but suceeds with the `Take[E, A]`.
+   * Error from stream when pulling is converted to `Take.halt`. Creates a singleton chunk.
+   */
+  def fromEffect[R, E, A](zio: ZIO[R, E, A]): URIO[R, Take[E, A]] =
+    zio.foldCause(halt, single)
+
+  /**
+   * Creates effect from `Pull[R, E, A]` that does not fail, but succeeds with the `Take[E, A]`.
+   * Error from stream when pulling is converted to `Take.halt`, end of stream to `Take.end`.
+   */
+  def fromPull[R, E, A](pull: ZStream.Pull[R, E, A]): URIO[R, Take[E, A]] =
+    pull.foldCause(Cause.sequenceCauseOption(_).fold[Take[E, Nothing]](end)(halt), chunk)
+
+  /**
+   * Creates a failing `Take[E, Nothing]` with the specified cause.
+   */
+  def halt[E](c: Cause[E]): Take[E, Nothing] =
+    Take(Exit.halt(c.map(Some(_))))
+
+  /**
+   * Creates a failing `Take[Nothing, Nothing]` with the specified throwable.
+   */
+  def die(t: Throwable): Take[Nothing, Nothing] =
+    Take(Exit.die(t))
+
+  /**
+   * Creates a failing `Take[Nothing, Nothing]` with the specified error message.
+   */
+  def dieMessage(msg: String): Take[Nothing, Nothing] =
+    Take(Exit.die(new RuntimeException(msg)))
+
+  /**
+   * Creates a `Take[E, A]` from `Exit[E, A]`.
+   */
+  def done[E, A](exit: Exit[E, A]): Take[E, A] =
+    Take(exit.mapError[Option[E]](Some(_)).map(Chunk.single))
+
+  /**
+   * End-of-stream marker
+   */
+  val end: Take[Nothing, Nothing] =
+    Take(Exit.fail(None))
+}