Mapping responses

6 files changed, 83 insertions, 43 deletions

diff --git a/README.md b/README.md
index 8c0de51..a64cfd3 100644
--- a/README.md
+++ b/README.md
@@ -104,8 +104,8 @@ val response: Response[String] = request.send()
 By default the response body is read into a utf-8 string. How the response body
 is handled is also part of the request description. The body can be ignore
 (`.response(ignore)`), read into a sequence of parameters 
-(`.response(asParams)`) and more; some backends also support request & response
-streaming.
+(`.response(asParams)`), mapped (`.mapResponse`) and more; some backends also 
+support request & response streaming.
 
 The default handler doesn't wrap the response into any container, but other
 asynchronous handlers might do so. The type parameter in the `Response[_]`
diff --git a/akka-http-handler/src/main/scala/com/softwaremill/sttp/akkahttp/AkkaHttpSttpHandler.scala b/akka-http-handler/src/main/scala/com/softwaremill/sttp/akkahttp/AkkaHttpSttpHandler.scala
index 62e066b..01f68ef 100644
--- a/akka-http-handler/src/main/scala/com/softwaremill/sttp/akkahttp/AkkaHttpSttpHandler.scala
+++ b/akka-http-handler/src/main/scala/com/softwaremill/sttp/akkahttp/AkkaHttpSttpHandler.scala
@@ -63,21 +63,21 @@ class AkkaHttpSttpHandler(actorSystem: ActorSystem)
       asByteArray.map(new String(_, enc))
 
     rr match {
-      case IgnoreResponse =>
+      case IgnoreResponse(g) =>
         hr.discardEntityBytes()
-        Future.successful(())
+        Future.successful(g(()))
 
-      case ResponseAsString(enc) =>
-        asString(enc)
+      case ResponseAsString(enc, g) =>
+        asString(enc).map(g)
 
-      case ResponseAsByteArray =>
-        asByteArray
+      case ResponseAsByteArray(g) =>
+        asByteArray.map(g)
 
-      case r @ ResponseAsParams(enc) =>
-        asString(enc).map(r.parse)
+      case r @ ResponseAsParams(enc, g) =>
+        asString(enc).map(r.parse).map(g)
 
-      case r @ ResponseAsStream() =>
-        Future.successful(r.responseIsStream(hr.entity.dataBytes))
+      case r @ ResponseAsStream(g) =>
+        Future.successful(r.responseIsStream(hr.entity.dataBytes)).map(g)
     }
   }
 
diff --git a/core/src/main/scala/com/softwaremill/sttp/HttpURLConnectionSttpHandler.scala b/core/src/main/scala/com/softwaremill/sttp/HttpURLConnectionSttpHandler.scala
index a42e26d..208f3a2 100644
--- a/core/src/main/scala/com/softwaremill/sttp/HttpURLConnectionSttpHandler.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/HttpURLConnectionSttpHandler.scala
@@ -93,15 +93,15 @@ object HttpURLConnectionSttpHandler extends SttpHandler[Id, Nothing] {
     def asString(enc: String) = Source.fromInputStream(is, enc).mkString
 
     responseAs match {
-      case IgnoreResponse =>
+      case IgnoreResponse(g) =>
         @tailrec def consume(): Unit = if (is.read() != -1) consume()
 
-        consume()
+        g(consume())
 
-      case ResponseAsString(enc) =>
-        asString(enc)
+      case ResponseAsString(enc, g) =>
+        g(asString(enc))
 
-      case ResponseAsByteArray =>
+      case ResponseAsByteArray(g) =>
         val os = new ByteArrayOutputStream
         var read = 0
         val buf = new Array[Byte](1024)
@@ -117,12 +117,12 @@ object HttpURLConnectionSttpHandler extends SttpHandler[Id, Nothing] {
 
         transfer()
 
-        os.toByteArray
+        g(os.toByteArray)
 
-      case r @ ResponseAsParams(enc) =>
-        r.parse(asString(enc))
+      case r @ ResponseAsParams(enc, g) =>
+        g(r.parse(asString(enc)))
 
-      case ResponseAsStream() =>
+      case ResponseAsStream(_) =>
         // only possible when the user requests the response as a stream of
         // Nothing. Oh well ...
         throw new IllegalStateException()
diff --git a/core/src/main/scala/com/softwaremill/sttp/model/package.scala b/core/src/main/scala/com/softwaremill/sttp/model/package.scala
index 6fb41b8..def6132 100644
--- a/core/src/main/scala/com/softwaremill/sttp/model/package.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/model/package.scala
@@ -46,18 +46,32 @@ package object model {
     * @tparam T Target type as which the response will be read.
     * @tparam S If `T` is a stream, the type of the stream. Otherwise, `Nothing`.
     */
-  sealed trait ResponseAs[T, +S]
+  sealed trait ResponseAs[T, +S] {
+    def map[T2](f: T => T2): ResponseAs[T2, S]
+  }
 
-  object IgnoreResponse extends ResponseAs[Unit, Nothing]
-  case class ResponseAsString(encoding: String)
-      extends ResponseAs[String, Nothing]
-  object ResponseAsByteArray extends ResponseAs[Array[Byte], Nothing]
-  case class ResponseAsStream[T, S]()(implicit val responseIsStream: S =:= T)
-      extends ResponseAs[T, S]
-  case class ResponseAsParams(encoding: String)
-      extends ResponseAs[Seq[(String, String)], Nothing] {
+  case class IgnoreResponse[T](g: Unit => T) extends ResponseAs[T, Nothing] {
+    override def map[T2](f: T => T2): ResponseAs[T2, Nothing] =
+      IgnoreResponse(g andThen f)
+  }
+  case class ResponseAsString[T](encoding: String, g: String => T)
+      extends ResponseAs[T, Nothing] {
+    override def map[T2](f: T => T2): ResponseAs[T2, Nothing] =
+      ResponseAsString(encoding, g andThen f)
+  }
+  case class ResponseAsByteArray[T](g: Array[Byte] => T)
+      extends ResponseAs[T, Nothing] {
+    override def map[T2](f: T => T2): ResponseAs[T2, Nothing] =
+      ResponseAsByteArray(g andThen f)
+  }
+  case class ResponseAsParams[T](encoding: String,
+                                 g: Seq[(String, String)] => T)
+      extends ResponseAs[T, Nothing] {
 
-    def parse(s: String): Seq[(String, String)] = {
+    override def map[T2](f: T => T2): ResponseAs[T2, Nothing] =
+      ResponseAsParams(encoding, g andThen f)
+
+    private[sttp] def parse(s: String): Seq[(String, String)] = {
       s.split("&")
         .toList
         .flatMap(kv =>
@@ -70,4 +84,11 @@ package object model {
         })
     }
   }
+  case class ResponseAsStream[T, T2, S](g: T => T2)(
+      implicit val responseIsStream: S =:= T)
+      extends ResponseAs[T2, S] {
+
+    override def map[T3](f: T2 => T3): ResponseAs[T3, S] =
+      ResponseAsStream(g andThen f)
+  }
 }
diff --git a/core/src/main/scala/com/softwaremill/sttp/package.scala b/core/src/main/scala/com/softwaremill/sttp/package.scala
index 925cb3b..d51d6c5 100644
--- a/core/src/main/scala/com/softwaremill/sttp/package.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/package.scala
@@ -16,16 +16,16 @@ package object sttp {
   type Id[X] = X
   type Empty[X] = None.type
 
-  def ignore: ResponseAs[Unit, Nothing] = IgnoreResponse
+  def ignore: ResponseAs[Unit, Nothing] = IgnoreResponse(identity)
 
   /**
     * Uses `utf-8` encoding.
     */
   def asString: ResponseAs[String, Nothing] = asString(Utf8)
   def asString(encoding: String): ResponseAs[String, Nothing] =
-    ResponseAsString(encoding)
+    ResponseAsString(encoding, identity)
   def asByteArray: ResponseAs[Array[Byte], Nothing] =
-    ResponseAsByteArray
+    ResponseAsByteArray(identity)
 
   /**
     * Uses `utf-8` encoding.
@@ -33,9 +33,9 @@ package object sttp {
   def asParams: ResponseAs[Seq[(String, String)], Nothing] =
     asParams(Utf8)
   def asParams(encoding: String): ResponseAs[Seq[(String, String)], Nothing] =
-    ResponseAsParams(encoding)
+    ResponseAsParams(encoding, identity)
 
-  def asStream[S]: ResponseAs[S, S] = ResponseAsStream[S, S]()
+  def asStream[S]: ResponseAs[S, S] = ResponseAsStream[S, S, S](identity)
 
   /**
     * Use the factory methods `multiPart` to conveniently create instances of
@@ -282,17 +282,18 @@ package object sttp {
       this.copy[U, T, S2](body = StreamBody(b))
 
     /**
-      * What's the target type to which the response body
-      *                   should be read. Needs to be specified upfront
-      *                   so that the response is always consumed and hence
-      *                   there are no requirements on client code to consume
-      *                   it. An exception to this are streaming responses,
-      *                   which need to fully consumed by the client if such
-      *                   a response type is requested.
+      * What's the target type to which the response body should be read.
+      * Needs to be specified upfront so that the response is always consumed
+      * and hence there are no requirements on client code to consume it. An
+      * exception to this are streaming responses, which need to fully
+      * consumed by the client if such a response type is requested.
       */
     def response[T2, S2 >: S](ra: ResponseAs[T2, S2]): RequestT[U, T2, S2] =
       this.copy(responseAs = ra)
 
+    def mapResponse[T2](f: T => T2): RequestT[U, T2, S] =
+      this.copy(responseAs = responseAs.map(f))
+
     def send[R[_]]()(implicit handler: SttpHandler[R, S],
                      isIdInRequest: IsIdInRequest[U]): R[Response[T]] = {
       // we could avoid the asInstanceOf by creating an artificial copy
diff --git a/tests/src/test/scala/com/softwaremill/sttp/BasicTests.scala b/tests/src/test/scala/com/softwaremill/sttp/BasicTests.scala
index 621968b..2b3bbab 100644
--- a/tests/src/test/scala/com/softwaremill/sttp/BasicTests.scala
+++ b/tests/src/test/scala/com/softwaremill/sttp/BasicTests.scala
@@ -137,6 +137,24 @@ class BasicTests
         response.body should be(expectedPostEchoResponse)
       }
 
+      name should "parse response as string with mapping using map" in {
+        val response = postEcho
+          .body(testBody)
+          .response(asString.map(_.length))
+          .send()
+          .force()
+        response.body should be(expectedPostEchoResponse.length)
+      }
+
+      name should "parse response as string with mapping using mapResponse" in {
+        val response = postEcho
+          .body(testBody)
+          .mapResponse(_.length)
+          .send()
+          .force()
+        response.body should be(expectedPostEchoResponse.length)
+      }
+
       name should "parse response as a byte array" in {
         val response =
           postEcho.body(testBody).response(asByteArray).send().force()