Play result composition

common-package.scala

/* Caution: untested code. Allows composing Result objects in Play 2.0 */


import play.api.mvc._
import play.api.http.HeaderNames

package object common {
    def mapPlain(result: Result, f: PlainResult => Result): Result = {
        result match {
            case p: PlainResult =>
                f(p)
            case a: AsyncResult =>
                AsyncResult(a.result.map(mapPlain(_, f)))
            case _ =>
                throw new IllegalArgumentException("unable to work with result: " + result)
        }
    }

    def mapHeaders(result: Result, f: ResponseHeader => ResponseHeader): Result = {
        mapPlain(result, { plain =>
            plain match {
                case simple: SimpleResult[_] =>
                    SimpleResult(f(simple.header), simple.body)(simple.writeable)
                case chunked: ChunkedResult[_] =>
                    ChunkedResult(f(chunked.header), chunked.chunks)(chunked.writeable)
                case _ =>
                    throw new IllegalArgumentException("unable to work with result: " + plain)
            }
        })
    }

    def mapCookies(result: Result, f: Seq[Cookie] => Seq[Cookie]): Result = {
        mapHeaders(result, { header =>
            val newCookies = f(header.headers.get(HeaderNames.SET_COOKIE).map(Cookies.decode(_)).getOrElse(Seq.empty))
            header.copy(headers = header.headers + (HeaderNames.SET_COOKIE -> Cookies.encode(newCookies)))
        })
    }

    def mapSession(result: Result, f: Session => Session): Result = {
        mapCookies(result, { cookies =>
            val (sessionCookies, otherCookies) = cookies.partition(_.name == Session.COOKIE_NAME)
            val oldSession = Session.decodeFromCookie(sessionCookies.headOption)
            val newSession = f(oldSession)
            Session.encodeAsCookie(newSession) +: otherCookies
        })
    }
}