Scala Seq distinct on an item property

seq-distinct-on.scala

  object SeqWithDistinctOn {
    /** Utility class to run a distinct function on */
    case class DistinctResult[A, B](val items: Seq[A] = Seq.empty, val keys: Seq[B] = Seq.empty) {
      /**
       * Add item with given key.
       * If key already exists, the item will not be added.
       *
       * @param item Item to be added if #key has not been added yet.
       * @param key Key of item.
       * @return A [[DistinctResult]] with added item of key did not exist, otherwise current [[DistinctResult]].
       */
      def addItemWithKey(item: A, key: B): DistinctResult[A, B] = if (keys contains key) this else copy(
        items = items :+ item,
        keys = keys :+ key)
    }
  }

  /**
   * Extend a [[Seq]] with a function on which to define on which item property
   * to do a 'distinct' function instead of on the whole item.
   */
  implicit class SeqWithDistinctOn[A](val seq: Seq[A]) extends AnyVal {
    import SeqWithDistinctOn.DistinctResult

    /**
     * Return a distinct Seq of items where each key determined by function #fKeyOfA only exists once.
     *
     * @param fKeyOfA Function to determine the unique key for an item.
     * @tparam B Type of item key.
     * @return Seq of unique items based on a key.
     */
    def distinctOn[B](fKeyOfA: A => B): Seq[A] = seq.foldLeft(DistinctResult[A, B]()) {
      case (result, item) => result.addItemWithKey(item, fKeyOfA(item))
    }.items
  }

import org.scalatest.{ Matchers, WordSpec }

class SeqWithDistinctOnSpec extends WordSpec with Matchers {

  case class Item(key: String)

  "SeqWithDistinctOn" should {
    import CollectionUtils.SeqWithDistinctOn

    "Distinct on item properties" in {

      val itemSeq = Seq(Item("one"), Item("two"), Item("three"), Item("four"), Item("five"))

      val distinctFirstChar = itemSeq.distinctOn(_.key.head)
      distinctFirstChar should contain theSameElementsInOrderAs (
        Seq(Item("one"), Item("two"), Item("four"))
      )
    }
  }
}