FocusedAgent

FocusedAgent.scala

import akka.agent.Agent
import shapeless.Lens
import concurrent.Future
import akka.util.Timeout
import concurrent.ExecutionContext

object FocusedAgent {
  implicit class EnrichedAgent[A](val agent: Agent[A]) extends AnyVal {
    def focusOn[B](lens: Lens[A,B]): FocusedAgent[A,B] = new FocusedAgent[A,B](agent, lens)
  }
}

class FocusedAgent[U,T](underlying: Agent[U], lens: Lens[U,T]) {
  def get(): T = lens get underlying.get()

  def apply(): T = get

  private def adaptedModify(f: T => T): U => U = lens.modify(_)(f)
  private def adaptedSet(x: T): U => U = lens.set(_)(x)

  def send(f: T ⇒ T): Unit = { underlying send adaptedModify(f) }

  def alter(f: T ⇒ T)(implicit timeout: Timeout, ec: ExecutionContext): Future[T] = {
    underlying alter adaptedModify(f) map lens.get
  }

  def send(newValue: T): Unit = underlying send adaptedSet(newValue)

  def update(newValue: T): Unit = send(newValue)

  def sendOff(f: T ⇒ T)(implicit ec: ExecutionContext): Unit = {
    underlying.sendOff(lens.modify(_)(f))
  }

  def alterOff(f: T ⇒ T)(implicit timeout: Timeout, ec: ExecutionContext): Future[T] = {
    underlying.alterOff(lens.modify(_)(f)) map lens.get
  }

  def future(implicit timeout: Timeout, ec: ExecutionContext): Future[T] =
    underlying.future map lens.get

  def await(implicit timeout: Timeout): T = lens get underlying.await

  def map[B](f: T ⇒ B): Agent[B] = underlying.map(f compose lens.get)

  def flatMap[B](f: T ⇒ Agent[B]): Agent[B] = f(get)

  def foreach[U](f: T ⇒ U): Unit = f(get)

  def suspend(): Unit = underlying.suspend()

  def resume(): Unit = underlying.resume()

  def close(): Unit = underlying.close()
}