Lenses in Scala with support for Maps, Seqs, and JSON (using the Jerkson library)

lenses.scala

package lenses

import com.codahale.jerkson._, AST._

/**
 * Lenses allow retrieving and updating values of type `B` inside some enclosing value of type `A`.
 *
 * Since lenses can be composed together, arbitrary nesting is automatically taken care of.
 */
case class Lens[A, B](get: A => B, set: (A, B) => A) {
  /**
   * Composes two lenses.
   */
  def compose[C](that: Lens[C, A]): Lens[C, B] = Lens(
    get = a => this.get(that.get(a)),
    set = (c, b) => {
      val a = that.get(c)
      that.set(c, this.set(a, b))
    })

  /**
   * Alias for `compose` with the parameters swapped.
   */
  def andThen[C](that: Lens[B, C]): Lens[A, C] = that compose this

  /**
   * A lens that transparently converts the value of type `B` to some other type `C`.
   */
  def withConversion[C](to: B => C)(from: C => B): Lens[A, C] = Lens(
    get = (a) => to(this.get(a)),
    set = (a, c) => this.set(a, from(c)))

  /**
   * Changes a lens that focuses on an optional value into a lens that replaces `None` with a default value.
   */
  def withDefault[C](default: C)(implicit ev1: Option[C] =:= B, ev2: B =:= Option[C]): Lens[A, C] = Lens(
    get = (a) => this.get(a).getOrElse(default),
    set = (a, c) => this.set(a, Some(c)))
}

object Lens {
  /**
   * A lens that simply focuses on itself.
   */
  def identity[A]: Lens[A, A] = Lens(get = (a) => a, set = (_, a) => a)

  /**
   * A lens that focuses on the specified key in a `Map`.
   */
  def forMap[K, V](key: K): Lens[Map[K, V], Option[V]] = Lens(
    get = (map) => map.get(key),
    set = (map, value) =>
      value match {
        case Some(value) => map.updated(key, value)
        case None        => map - key
      })

  /**
   * A lens that focuses on the specified index in a `Seq`.
   */
  def forSeq[V, CC <: collection.GenSeqLike[V, CC]](index: Int)(implicit bf: collection.generic.CanBuildFrom[CC, V, CC]): Lens[CC, V] = Lens(
    get = (seq) => seq(index),
    set = (seq, value) => seq.updated(index, value))

  /**
   * A lens that focuses on a specific field inside a JSON object.
   */
  def forJsonField(key: String): Lens[JValue, JValue] = Lens(
    get = json => json \ key,
    set = {
      case (JNull, JNull) =>
        JNull
      case (JNull, value) =>
        JObject(List(JField(key, value)))
      case (JObject(fields), JNull) =>
        JObject(fields.filter(_.name != key))
      case (JObject(fields), value) =>
        fields.span(_.name != key) match {
          case (prefix, _ :: suffix) => JObject(prefix ::: JField(key, value) :: suffix)
          case (_, Nil)              => JObject(JField(key, value) :: fields)
        }
      case (json, _) =>
        json
    })

  /**
   * A lens that focuses on a specific element inside a JSON array.
   */
  def forJsonArray(index: Int): Lens[JValue, JValue] = Lens(
    get = json => json(index),
    set = {
      case (JArray(elements), value) if elements.indices contains index =>
        JArray(elements.updated(index, value))
      case (json, _) =>
        json
    })
}

object LensExamples {
  case class Invoice(organization: Organization, number: String, totalAmount: Double)
  case class Organization(name: String, attributes: JValue)

  val attributesL: Lens[Organization, JValue] = Lens(get = _.attributes, set = (organization, attributes) => organization.copy(attributes = attributes))
  val organizationL: Lens[Invoice, Organization] = Lens(get = _.organization, set = (invoice, organization) => invoice.copy(organization = organization))

  val healthcareAttributesL: Lens[JValue, JValue] = Lens.forJsonField("healthcare")
  val bigNumberAttributeL: Lens[JValue, Option[String]] = Lens.forJsonField("bigNumber").withConversion {
    case JString(value) => Some(value)
    case JNull          => None
    case x              => throw new MatchError("cannot convert " + x + " to a String")
  } {
    case Some(value) => JString(value)
    case None        => JNull
  }

  val invoiceBigNumberL = organizationL andThen attributesL andThen healthcareAttributesL andThen bigNumberAttributeL

  val example1 = Invoice(
    number = "1001",
    totalAmount = 2.71,
    organization = Organization(name = "Bar BV", attributes = JNull))
  val example2 = Invoice(
    number = "1234",
    totalAmount = 3.14,
    organization = Organization(name = "Foo BV", attributes = Json.parse[JValue]("""{"healthcare":{"bigNumber":"4242-123"}}""")))
}

repl.scala

scala> import com.codahale.jerkson._, AST._
import com.codahale.jerkson._
import AST._

scala> import lenses._, LensExamples._
import lenses._
import LensExamples._

scala> invoiceBigNumberL.get(example1)
res0: Option[String] = None

scala> invoiceBigNumberL.get(example2)
res1: Option[String] = Some(4242-123)

scala> example1
res2: lenses.LensExamples.Invoice = Invoice(Organization(Bar BV,JNull),1001,2.71)

scala> invoiceBigNumberL.set(example1, Some("3434-123"))
res3: lenses.LensExamples.Invoice = Invoice(Organization(Bar BV,JObject(List(JField(healthcare,JObject(List(JField(bigNumber,JString(3434-123)))))))),1001,2.71)
// ^^^ intermediate nested JSON is automatically created!

scala> invoiceBigNumberL.set(example2, Some("9999-999"))
res4: lenses.LensExamples.Invoice = Invoice(Organization(Foo BV,JObject(List(JField(healthcare,JObject(List(JField(bigNumber,JString(9999-999)))))))),1234,3.14)

scala> invoiceBigNumberL.set(example2, None)
res5: lenses.LensExamples.Invoice = Invoice(Organization(Foo BV,JObject(List(JField(healthcare,JObject(List()))))),1234,3.14)