Skip to content

Commit

Permalink
Add groupBy for NonEmptySet
Browse files Browse the repository at this point in the history
  • Loading branch information
Obarros committed Jun 12, 2018
1 parent 3993dba commit 2ac680e
Show file tree
Hide file tree
Showing 5 changed files with 74 additions and 0 deletions.
15 changes: 15 additions & 0 deletions core/src/main/scala/cats/data/NonEmptySet.scala
Original file line number Diff line number Diff line change
Expand Up @@ -346,6 +346,21 @@ sealed class NonEmptySetOps[A](val value: NonEmptySet[A]) {
*/
def zipWithIndex: NonEmptySet[(A, Int)] =
NonEmptySetImpl.create(toSortedSet.zipWithIndex)

/**
* Groups elements inside this `NonEmptySet` according to the `Order`
* of the keys produced by the given mapping function.
*/
def groupBy[B](f: A => B)(implicit B: Order[B]): NonEmptyMap[B, NonEmptySet[A]] = {
reduceLeftTo(a => NonEmptyMap.one(f(a), NonEmptySet.one(a))) { (acc, a) =>
val key = f(a)
val result = acc.lookup(key) match {
case Some(nes) => nes.add(a)
case _ => NonEmptySet.one(a)
}
acc.add((key, result))
}
}
}

private[data] sealed abstract class NonEmptySetInstances {
Expand Down
1 change: 1 addition & 0 deletions core/src/main/scala/cats/syntax/all.scala
Original file line number Diff line number Diff line change
Expand Up @@ -64,3 +64,4 @@ trait AllSyntaxBinCompat1
with NestedSyntax
with BinestedSyntax
with ParallelFlatSyntax
with SetSyntax
1 change: 1 addition & 0 deletions core/src/main/scala/cats/syntax/package.scala
Original file line number Diff line number Diff line change
Expand Up @@ -52,4 +52,5 @@ package object syntax {
object validated extends ValidatedSyntax
object vector extends VectorSyntax
object writer extends WriterSyntax
object set extends SetSyntax
}
52 changes: 52 additions & 0 deletions core/src/main/scala/cats/syntax/set.scala
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
package cats.syntax

import scala.collection.immutable.{SortedSet, SortedMap}
import cats.data.NonEmptySet
import cats.Order

trait SetSyntax {
implicit final def catsSyntaxSet[A](se: SortedSet[A]): SetOps[A] = new SetOps(se)
}

final class SetOps[A](val se: SortedSet[A]) extends AnyVal {

/**
* Returns an Option of NonEmptySet from a SortedSet
*
* Example:
* {{{
* scala> import scala.collection.immutable.SortedSet
* scala> import cats.data.NonEmptySet
* scala> import cats.implicits._
*
* scala> val result1: SortedSet[Int] = SortedSet(1, 2)
* scala> result1.toNes
* res0: Option[NonEmptySet[Int]] = Some(TreeSet(1, 2))
*
* scala> val result2: SortedSet[Int] = SortedSet.empty[Int]
* scala> result2.toNes
* res1: Option[NonEmptySet[Int]] = None
* }}}
*/
def toNes: Option[NonEmptySet[A]] = NonEmptySet.fromSet(se)

/**
* Groups elements inside this `SortedSet` according to the `Order` of the keys
* produced by the given mapping function.
*
* {{{
* scala> import cats.data.NonEmptySet
* scala> import scala.collection.immutable.{SortedMap, SortedSet}
* scala> import cats.implicits._
*
* scala> val sortedSet = SortedSet(12, -2, 3, -5)
*
* scala> sortedSet.groupByNes(_ >= 0)
* res0: SortedMap[Boolean, NonEmptySet[Int]] = Map(false -> TreeSet(-5, -2), true -> TreeSet(3, 12))
* }}}
*/
def groupByNes[B](f: A => B)(implicit B: Order[B]): SortedMap[B, NonEmptySet[A]] = {
implicit val ordering = B.toOrdering
toNes.fold(SortedMap.empty[B, NonEmptySet[A]])(_.groupBy(f).toSortedMap)
}
}
5 changes: 5 additions & 0 deletions tests/src/test/scala/cats/tests/NonEmptySetSuite.scala
Original file line number Diff line number Diff line change
Expand Up @@ -233,4 +233,9 @@ class NonEmptySetSuite extends CatsSuite {
}
}

test("NonEmptySet#groupBy is consistent with Set#groupBy") {
forAll { (nes: NonEmptySet[Int], f: Int => Int) =>
nes.groupBy(f).map(_.toSortedSet).toSortedMap should === (nes.toSortedSet.groupBy(f))
}
}
}

0 comments on commit 2ac680e

Please sign in to comment.