Add context method to request context (#2799)

* Add `context` method to request context

* Remove dead code

zio-http/jvm/src/test/scala/zio/http/HandlerAspectSpec.scala

@@ -18,7 +18,7 @@ object HandlerAspectSpec extends ZIOSpecDefault {
       test("HandlerAspect with context can eliminate environment type partially") {
         val handlerAspect = HandlerAspect.interceptIncomingHandler(handler((req: Request) => (req, req.headers.size)))
         val handler0 = handler { (_: Request) =>
-          ZIO.service[Boolean] *> ZIO.serviceWith[Int](i => Response.text(i.toString))
+          withContext((_: Boolean, i: Int) => Response.text(i.toString))
           //leftover type is only needed in Scala 2
           //can't be infix because of Scala 3
         }.@@[Boolean](handlerAspect)
@@ -29,10 +29,10 @@ object HandlerAspectSpec extends ZIOSpecDefault {
       },
       test("HandlerAspect with context can eliminate environment type partially while requiring an additional environment") {
         val handlerAspect: HandlerAspect[String, Int] = HandlerAspect.interceptIncomingHandler {
-          handler((req: Request) => ZIO.serviceWith[String](s => (req.withBody(Body.fromString(s)), req.headers.size)))
+          handler((req: Request) => withContext((s: String) => (req.withBody(Body.fromString(s)), req.headers.size)))
         }
         val handler0: Handler[Boolean with String, Response, Request, Response] = handler { (r: Request) =>
-          ZIO.service[Boolean] *> ZIO.serviceWithZIO[Int] { i =>
+          ZIO.service[Boolean] *> withContext{ (i: Int) =>
             for {
               body <- r.body.asString.orDie
             } yield Response.text(s"$i $body")

zio-http/jvm/src/test/scala/zio/http/internal/middlewares/AuthSpec.scala

@@ -19,7 +19,7 @@ package zio.http.internal.middlewares
 import zio.Config.Secret
 import zio.test.Assertion._
 import zio.test._
-import zio.{Ref, ZEnvironment, ZIO}
+import zio.{Ref, ZIO}
 
 import zio.http._
 import zio.http.internal.HttpAppTestExtensions
@@ -74,7 +74,7 @@ object AuthSpec extends ZIOHttpSpec with HttpAppTestExtensions {
       },
       test("Extract username via context") {
         val app = (Handler.fromFunctionZIO[Request](_ =>
-          ZIO.serviceWith[AuthContext](c => Response.text(c.value)),
+          withContext((c: AuthContext) => Response.text(c.value)),
         ) @@ basicAuthContextM).merge.mapZIO(_.body.asString)
         assertZIO(app.runZIO(Request.get(URL.empty).copy(headers = successBasicHeader)))(equalTo("user"))
       },
@@ -109,14 +109,13 @@ object AuthSpec extends ZIOHttpSpec with HttpAppTestExtensions {
         val secureRoutes = Routes(
           Method.GET / "a" -> handler((_: Request) => ZIO.serviceWith[AuthContext](ctx => Response.text(ctx.value))),
           Method.GET / "b" / int("id")      -> handler((id: Int, _: Request) =>
-            ZIO.serviceWith[AuthContext](ctx => Response.text(s"for id: $id: ${ctx.value}")),
+            withContext((ctx: AuthContext) => Response.text(s"for id: $id: ${ctx.value}")),
           ),
           Method.GET / "c" / string("name") -> handler((name: String, _: Request) =>
-            ZIO.serviceWith[AuthContext](ctx => Response.text(s"for name: $name: ${ctx.value}")),
+            withContext((ctx: AuthContext) => Response.text(s"for name: $name: ${ctx.value}")),
           ),
           // Needs version of @@ that removes the context from the environment
         ) @@ basicAuthContextM
-        // Just a prove that the aspect can require an environment. Does nothing.
         val app          = secureRoutes
         for {
           s1     <- app.runZIO(Request.get(URL(Path.root / "a")).copy(headers = successBasicHeader))

zio-http/shared/src/main/scala/zio/http/WithContext.scala

@@ -0,0 +1,316 @@
+/*
+ * Copyright 2023 the ZIO HTTP 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.http
+
+import scala.annotation.implicitNotFound
+
+import zio._
+
+@implicitNotFound("""
+The type ${C} does not appear to be one that can be used when requesting a context. The following types may be used when requesting a context:
+
+ - Simple Functions:  Ctx => Out
+ - ZIO Functions:     Ctx => ZIO[Env, Err, Out]
+
+The types to the function must be specified explicitly: Scala cannot infer them due to the smart constructor.
+""")
+trait WithContext[C] {
+  type Env
+  type Err
+  type Out
+
+  def toZIO(h: => C): ZIO[Env, Err, Out]
+}
+
+object WithContext extends WithContextConstructorLowPriorityImplicits1
+
+private[http] trait WithContextConstructorLowPriorityImplicits1 extends WithContextConstructorLowPriorityImplicits2 {
+
+  implicit def function2ZIOWithContextConstructor[Env0, Err0, Ctx1, Ctx2, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+  ): WithContext.Typed[(Ctx1, Ctx2) => ZIO[Env0, Err0, Out0], Env0 with Ctx1 with Ctx2, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2) => ZIO[Env0, Err0, Out0]] {
+      type Env = Env0 with Ctx1 with Ctx2
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2) => ZIO[Env0, Err0, Out0]
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] = {
+        implicit val tag1_ = tag1.tag
+        implicit val tag2_ = tag2.tag
+        implicit val usf   = Unsafe.unsafe
+
+        ZIO.suspendSucceed {
+          FiberRef.currentEnvironment.get.flatMap(environment =>
+            z(environment.unsafe.get(tag1_), environment.unsafe.get(tag2_)),
+          )
+        }
+      }
+
+    }
+
+  implicit def function3ZIOWithContextConstructor[Env0, Err0, Ctx1, Ctx2, Ctx3, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+    tag3: Tag[Ctx3],
+  ): WithContext.Typed[(Ctx1, Ctx2, Ctx3) => ZIO[Env0, Err0, Out0], Env0 with Ctx1 with Ctx2 with Ctx3, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2, Ctx3) => ZIO[Env0, Err0, Out0]] {
+      type Env = Env0 with Ctx1 with Ctx2 with Ctx3
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2, Ctx3) => ZIO[Env0, Err0, Out0]
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val tag3_ = tag3.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.flatMap(environment =>
+            z(environment.unsafe.get(tag1_), environment.unsafe.get(tag2_), environment.unsafe.get(tag3_)),
+          )
+        }
+
+    }
+
+  implicit def function4ZIOWithContextConstructor[Env0, Err0, Ctx1, Ctx2, Ctx3, Ctx4, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+    tag3: Tag[Ctx3],
+    tag4: Tag[Ctx4],
+  ): WithContext.Typed[(Ctx1, Ctx2, Ctx3, Ctx4) => ZIO[
+    Env0,
+    Err0,
+    Out0,
+  ], Env0 with Ctx1 with Ctx2 with Ctx3 with Ctx4, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2, Ctx3, Ctx4) => ZIO[Env0, Err0, Out0]] {
+      type Env = Env0 with Ctx1 with Ctx2 with Ctx3 with Ctx4
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2, Ctx3, Ctx4) => ZIO[Env0, Err0, Out0]
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val tag3_ = tag3.tag
+          implicit val tag4_ = tag4.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.flatMap(environment =>
+            z(
+              environment.unsafe.get(tag1_),
+              environment.unsafe.get(tag2_),
+              environment.unsafe.get(tag3_),
+              environment.unsafe.get(tag4_),
+            ),
+          )
+        }
+
+    }
+
+  implicit def function5ZIOWithContextConstructor[Env0, Err0, Ctx1, Ctx2, Ctx3, Ctx4, Ctx5, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+    tag3: Tag[Ctx3],
+    tag4: Tag[Ctx4],
+    tag5: Tag[Ctx5],
+  ): WithContext.Typed[(Ctx1, Ctx2, Ctx3, Ctx4, Ctx5) => ZIO[
+    Env0,
+    Err0,
+    Out0,
+  ], Env0 with Ctx1 with Ctx2 with Ctx3 with Ctx4 with Ctx5, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2, Ctx3, Ctx4, Ctx5) => ZIO[Env0, Err0, Out0]] {
+      type Env = Env0 with Ctx1 with Ctx2 with Ctx3 with Ctx4 with Ctx5
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2, Ctx3, Ctx4, Ctx5) => ZIO[Env0, Err0, Out0]
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val tag3_ = tag3.tag
+          implicit val tag4_ = tag4.tag
+          implicit val tag5_ = tag5.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.flatMap(environment =>
+            z(
+              environment.unsafe.get(tag1_),
+              environment.unsafe.get(tag2_),
+              environment.unsafe.get(tag3_),
+              environment.unsafe.get(tag4_),
+              environment.unsafe.get(tag5_),
+            ),
+          )
+        }
+    }
+
+}
+
+private[http] trait WithContextConstructorLowPriorityImplicits2 extends WithContextConstructorLowPriorityImplicits3 {
+  implicit def functionZIOWithContextConstructor[Env0, Err0, Ctx0, Out0](implicit
+    tag: Tag[Ctx0],
+  ): WithContext.Typed[Ctx0 => ZIO[Env0, Err0, Out0], Env0 with Ctx0, Err0, Out0] =
+    new WithContext[Ctx0 => ZIO[Env0, Err0, Out0]] {
+      type Env = Env0 with Ctx0
+      type Err = Err0
+      type Out = Out0
+      type Z   = Ctx0 => ZIO[Env0, Err0, Out0]
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] = {
+        implicit val tag_ = tag.tag
+        implicit val usf  = Unsafe.unsafe
+
+        ZIO.suspendSucceed {
+          FiberRef.currentEnvironment.get.flatMap(environment => z(environment.unsafe.get(tag_)))
+        }
+      }
+    }
+}
+
+private[http] trait WithContextConstructorLowPriorityImplicits3 extends WithContextConstructorLowPriorityImplicits4 {
+  implicit def function2ValueWithContextConstructor[Err0, Ctx1, Ctx2, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+  ): WithContext.Typed[(Ctx1, Ctx2) => Out0, Ctx1 with Ctx2, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2) => Out0] {
+      type Env = Ctx1 with Ctx2
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2) => Out0
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.map(environment =>
+            z(environment.unsafe.get(tag1_), environment.unsafe.get(tag2_)),
+          )
+        }
+
+    }
+
+  implicit def function3ValueWithContextConstructor[Err0, Ctx1, Ctx2, Ctx3, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+    tag3: Tag[Ctx3],
+  ): WithContext.Typed[(Ctx1, Ctx2, Ctx3) => Out0, Ctx1 with Ctx2 with Ctx3, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2, Ctx3) => Out0] {
+      type Env = Ctx1 with Ctx2 with Ctx3
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2, Ctx3) => Out0
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val tag3_ = tag3.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.map(environment =>
+            z(environment.unsafe.get(tag1_), environment.unsafe.get(tag2_), environment.unsafe.get(tag3_)),
+          )
+        }
+    }
+
+  implicit def function4ValueWithContextConstructor[Err0, Ctx1, Ctx2, Ctx3, Ctx4, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+    tag3: Tag[Ctx3],
+    tag4: Tag[Ctx4],
+  ): WithContext.Typed[(Ctx1, Ctx2, Ctx3, Ctx4) => Out0, Ctx1 with Ctx2 with Ctx3 with Ctx4, Err0, Out0] =
+    new WithContext[(Ctx1, Ctx2, Ctx3, Ctx4) => Out0] {
+      type Env = Ctx1 with Ctx2 with Ctx3 with Ctx4
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2, Ctx3, Ctx4) => Out0
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val tag3_ = tag3.tag
+          implicit val tag4_ = tag4.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.map(environment =>
+            z(
+              environment.unsafe.get(tag1_),
+              environment.unsafe.get(tag2_),
+              environment.unsafe.get(tag3_),
+              environment.unsafe.get(tag4_),
+            ),
+          )
+        }
+    }
+
+  implicit def function5ValueWithContextConstructor[Err0, Ctx1, Ctx2, Ctx3, Ctx4, Ctx5, Out0](implicit
+    tag1: Tag[Ctx1],
+    tag2: Tag[Ctx2],
+    tag3: Tag[Ctx3],
+    tag4: Tag[Ctx4],
+    tag5: Tag[Ctx5],
+  ): WithContext.Typed[
+    (Ctx1, Ctx2, Ctx3, Ctx4, Ctx5) => Out0,
+    Ctx1 with Ctx2 with Ctx3 with Ctx4 with Ctx5,
+    Err0,
+    Out0,
+  ] =
+    new WithContext[(Ctx1, Ctx2, Ctx3, Ctx4, Ctx5) => Out0] {
+      type Env = Ctx1 with Ctx2 with Ctx3 with Ctx4 with Ctx5
+      type Err = Err0
+      type Out = Out0
+      type Z   = (Ctx1, Ctx2, Ctx3, Ctx4, Ctx5) => Out0
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.suspendSucceed {
+          implicit val tag1_ = tag1.tag
+          implicit val tag2_ = tag2.tag
+          implicit val tag3_ = tag3.tag
+          implicit val tag4_ = tag4.tag
+          implicit val tag5_ = tag5.tag
+          implicit val usf   = Unsafe.unsafe
+          FiberRef.currentEnvironment.get.map(environment =>
+            z(
+              environment.unsafe.get(tag1_),
+              environment.unsafe.get(tag2_),
+              environment.unsafe.get(tag3_),
+              environment.unsafe.get(tag4_),
+              environment.unsafe.get(tag5_),
+            ),
+          )
+        }
+    }
+}
+
+private[http] trait WithContextConstructorLowPriorityImplicits4 {
+  type Typed[C, Env0, Err0, Out0] = WithContext[C] {
+    type Env = Env0; type Err = Err0; type Out = Out0
+  }
+
+  implicit def functionWithContextConstructor[Ctx: Tag, Out0]: WithContext.Typed[Ctx => Out0, Ctx, Nothing, Out0] =
+    new WithContext[Ctx => Out0] {
+      type Env = Ctx
+      type Err = Nothing
+      type Out = Out0
+      type Z   = Ctx => Out0
+
+      def toZIO(z: => Z): ZIO[Env, Err, Out] =
+        ZIO.serviceWith[Ctx] { in1 => z(in1) }
+    }
+}

zio-http/shared/src/main/scala/zio/http/package.scala

@@ -33,6 +33,9 @@ package object http extends UrlInterpolator with MdInterpolator {
   def handlerTODO(message: String): Handler[Any, Nothing, Any, Nothing] =
     handler(ZIO.dieMessage(message))
 
+  def withContext[C](fn: => C)(implicit c: WithContext[C]): ZIO[c.Env, c.Err, c.Out] =
+    c.toZIO(fn)
+
   abstract class RouteDecode[A](f: String => A) {
     def unapply(a: String): Option[A] =
       try {