Enumeratee.mapWithCounter

Enumeratee.mapWithCounter.scala

/**
 * Created by atamborrino on 11/03/2014.
 */

import play.api.libs.iteratee.Enumeratee._
import play.api.libs.iteratee._
import scala.concurrent._
import scala.concurrent.duration._

def mapWithCounter[From, To](f: (From, Int) => To): Enumeratee[From, To] = new CheckDone[From, To] {
  def step[A](counter: Int)(k: K[To, A]): K[From, Iteratee[To, A]] = {
    case in @ Input.El(_) =>
      new CheckDone[From, To] { def continue[A](k: K[To, A]) = Cont(step(counter + 1)(k)) } &> k(in.map(f(_, counter)))
    case in @ Input.Empty =>
      new CheckDone[From, To] { def continue[A](k: K[To, A]) = Cont(step(counter)(k)) } &> k(in)

    case Input.EOF => Done(Cont(k), Input.EOF)
  }

  def continue[A](k: K[To, A]) = Cont(step(0)(k))

}

val enum = Enumerator("a", "b") andThen Enumerator.empty andThen Enumerator("c")

val fut = enum &> mapWithCounter((e, counter) => (e, counter)) |>>> Iteratee.getChunks

Await.result(fut, 1 second) // res0: List[(String, Int)] = List((a,0), (b,1), (c,2))