some improvements (#276)

Add Logging.any #242
locally for ZStream #267
locally for Managed #210

build.sbt

@@ -51,6 +51,7 @@ lazy val core    = crossProject(JSPlatform, JVMPlatform)
   .settings(
     libraryDependencies ++= Seq(
       "dev.zio"                 %%% "zio"                     % ZioVersion,
+      "dev.zio"                 %%% "zio-streams"             % ZioVersion,
       ("org.scala-lang.modules" %%% "scala-collection-compat" % "2.4.4").cross(CrossVersion.for3Use2_13),
       "dev.zio"                 %%% "zio-test"                % ZioVersion % Test,
       "dev.zio"                 %%% "zio-test-sbt"            % ZioVersion % Test

core/shared/src/main/scala/zio/logging/Logger.scala

@@ -1,6 +1,7 @@
 package zio.logging
 
-import zio.{ Cause, FiberRef, UIO, URIO, ZIO }
+import zio.stream.ZStream
+import zio.{ Cause, FiberRef, UIO, URIO, ZIO, ZManaged }
 
 trait Logger[-A] { self =>
 
@@ -109,6 +110,18 @@ trait Logger[-A] { self =>
   def locallyM[R1, E, A1](f: LogContext => URIO[R1, LogContext])(zio: ZIO[R1, E, A1]): ZIO[R1, E, A1] =
     logContext.flatMap(ctx => f(ctx)).flatMap(ctx => locally(_ => ctx)(zio))
 
+  /**
+   * Modify log context in scope of Managed operation.
+   */
+  def locallyManaged[R1, E, A1](f: LogContext => LogContext)(managed: ZManaged[R1, E, A1]): ZManaged[R1, E, A1] =
+    ZManaged.makeReserve(managed.reserve.map(r => r.copy(locally(f)(r.acquire), exit => locally(f)(r.release(exit)))))
+
+  /**
+   * Modify log context in scope of ZStream.
+   */
+  def locallyZStream[R1, E, A1](f: LogContext => LogContext)(stream: ZStream[R1, E, A1]): ZStream[R1, E, A1] =
+    ZStream(stream.process.map(p => locally(f)(p)))
+
   /**
    * Modifies the annotate in the scope of the specified effect.
    */

core/shared/src/main/scala/zio/logging/Logging.scala

@@ -4,6 +4,7 @@ import zio._
 import zio.clock._
 import zio.console.Console
 import zio.logging.Logger.LoggerWithFormat
+import zio.stream.ZStream
 
 import java.nio.charset.{ Charset, StandardCharsets }
 import java.nio.file.Path
@@ -16,6 +17,8 @@ object Logging {
   def addTimestamp[A](self: Logger[A]): URIO[Clock, Logger[A]] =
     self.deriveM(ctx => currentDateTime.orDie.map(time => LogAnnotation.Timestamp(time)(ctx)))
 
+  val any: ZLayer[Logging, Nothing, Logging] = ZLayer.requires[Logging]
+
   def console(
     logLevel: LogLevel = LogLevel.Info,
     format: LogFormat[String] = LogFormat.ColoredLogFormat((_, s) => s)
@@ -97,6 +100,16 @@ object Logging {
   )(zio: ZIO[R, E, A1]): ZIO[Logging with R, E, A1] =
     ZIO.accessM(_.get.locallyM(fn)(zio))
 
+  def locallyManaged[A, R <: Logging, E, A1](fn: LogContext => LogContext)(
+    zio: ZManaged[R, E, A1]
+  ): ZManaged[Logging with R, E, A1] =
+    ZManaged.accessManaged(_.get.locallyManaged(fn)(zio))
+
+  def locallyZStream[A, R <: Logging, E, A1](fn: LogContext => LogContext)(
+    stream: ZStream[R, E, A1]
+  ): ZStream[Logging with R, E, A1] =
+    ZStream.accessStream(_.get.locallyZStream(fn)(stream))
+
   def make: URLayer[Appender[String], Logging] =
     ZLayer.fromFunctionM((appender: Appender[String]) =>
       FiberRef

core/shared/src/main/scala/zio/logging/log.scala

@@ -1,6 +1,7 @@
 package zio.logging
 
-import zio.{ Cause, URIO, ZIO }
+import zio.stream.ZStream
+import zio.{ Cause, URIO, ZIO, ZManaged }
 
 object log {
 
@@ -28,11 +29,21 @@ object log {
   def locally[R <: Logging, E, A1](fn: LogContext => LogContext)(zio: ZIO[R, E, A1]): ZIO[Logging with R, E, A1] =
     Logging.locally(fn)(zio)
 
+  def locallyManaged[R <: Logging, E, A1](fn: LogContext => LogContext)(
+    managed: ZManaged[R, E, A1]
+  ): ZManaged[Logging with R, E, A1] =
+    Logging.locallyManaged(fn)(managed)
+
   def locallyM[R <: Logging, E, A1](
     fn: LogContext => URIO[R, LogContext]
   )(zio: ZIO[R, E, A1]): ZIO[Logging with R, E, A1] =
     Logging.locallyM(fn)(zio)
 
+  def locallyZStream[R <: Logging, E, A1](fn: LogContext => LogContext)(
+    stream: ZStream[R, E, A1]
+  ): ZStream[Logging with R, E, A1] =
+    Logging.locallyZStream(fn)(stream)
+
   def throwable(line: => String, t: Throwable): ZIO[Logging, Nothing, Unit] =
     Logging.throwable(line, t)
 

core/shared/src/test/scala/zio/logging/LoggerSpec.scala

@@ -1,9 +1,10 @@
 package zio.logging
 
+import zio.stream.ZStream
 import zio.test.Assertion._
 import zio.test._
 import zio.test.environment.TestEnvironment
-import zio.{ FiberRef, Has, Layer, Ref, UIO, ZIO, ZLayer }
+import zio.{ FiberRef, Has, Layer, Ref, UIO, ZIO, ZLayer, ZManaged }
 
 import java.time.{ DateTimeException, OffsetDateTime }
 import java.util.UUID
@@ -139,6 +140,66 @@ object LoggerSpec extends DefaultRunnableSpec {
               )
             )
       },
+      testM("locally for Managed") {
+
+        log
+          .locallyManaged(LogAnnotation.Name("level-1" :: Nil)) {
+            ZManaged.make(log.info("acquire"))(_ => log.info("release"))
+          }
+          .use(_ => log.info("use")) *>
+          assertM(TestLogger.lines)(
+            equalTo(
+              Vector(
+                (
+                  LogContext.empty
+                    .annotate(LogAnnotation.Name, "level-1" :: Nil)
+                    .annotate(LogAnnotation.Level, LogLevel.Info),
+                  "acquire"
+                ),
+                (
+                  LogContext.empty
+                    .annotate(LogAnnotation.Level, LogLevel.Info),
+                  "use"
+                ),
+                (
+                  LogContext.empty
+                    .annotate(LogAnnotation.Name, "level-1" :: Nil)
+                    .annotate(LogAnnotation.Level, LogLevel.Info),
+                  "release"
+                )
+              )
+            )
+          )
+
+      },
+      testM("locally for Stream") {
+
+        log
+          .locallyZStream(LogAnnotation.Name("level-1" :: Nil)) {
+            ZStream.fromEffect(log.info("line1")) *>
+              ZStream.fromEffect(log.info("line2"))
+          }
+          .runDrain *>
+          assertM(TestLogger.lines)(
+            equalTo(
+              Vector(
+                (
+                  LogContext.empty
+                    .annotate(LogAnnotation.Name, "level-1" :: Nil)
+                    .annotate(LogAnnotation.Level, LogLevel.Info),
+                  "line1"
+                ),
+                (
+                  LogContext.empty
+                    .annotate(LogAnnotation.Name, "level-1" :: Nil)
+                    .annotate(LogAnnotation.Level, LogLevel.Info),
+                  "line2"
+                )
+              )
+            )
+          )
+
+      },
       test("LogContext render") {
         val correlationId = UUID.randomUUID()
         val rendered      = LogContext.empty

examples/src/main/scala/zio/logging/ZStreamExample.scala

@@ -0,0 +1,47 @@
+package zio.logging
+
+import zio._
+import zio.clock.Clock
+import zio.stream.ZStream
+
+import java.util.UUID
+
+object ZStreamExample extends zio.App {
+
+  final val CalculationId: LogAnnotation[Option[UUID]] = LogAnnotation[Option[UUID]](
+    name = "calculation-id",
+    initialValue = None,
+    combine = (_, r) => r,
+    render = _.map(_.toString).getOrElse("undefined-calculation-id")
+  )
+
+  final val CalculationNumber: LogAnnotation[Int] = LogAnnotation[Int](
+    name = "calculation-number",
+    initialValue = 0,
+    combine = (_, r) => r,
+    render = _.toString
+  )
+
+  final val env: ZLayer[zio.console.Console with Clock, Nothing, Logging] =
+    Logging.console(
+      logLevel = LogLevel.Debug,
+      format = LogFormat.ColoredLogFormat((ctx, line) => s"${ctx(CalculationId)} ${ctx(CalculationNumber)} $line")
+    ) >>>
+      Logging.withRootLoggerName("my-logger")
+
+  override def run(args: List[String]): URIO[zio.ZEnv, ExitCode] = {
+
+    val stream = for {
+      calcNumber <- ZStream(1, 2, 3, 4, 5)
+
+      subStream <-
+        log.locallyZStream(CalculationId(Some(UUID.randomUUID())).andThen(CalculationNumber(calcNumber)))(
+          ZStream.fromEffect(log.debug(s"would log first line for calculation# ${calcNumber}")) *>
+            ZStream.fromEffect(log.debug(s"would log second line for calculation# ${calcNumber}"))
+        )
+
+    } yield subStream
+
+    stream.runDrain.provideSomeLayer(env).exitCode
+  }
+}