From 7ce57d89e654e0ccac425e29a59f5cc176a48a81 Mon Sep 17 00:00:00 2001 From: obarros Date: Sat, 9 Jun 2018 12:33:04 +0100 Subject: [PATCH] Add groupBy for NonEmptySet --- .../main/scala/cats/data/NonEmptySet.scala | 15 ++++++ core/src/main/scala/cats/syntax/all.scala | 1 + core/src/main/scala/cats/syntax/package.scala | 1 + core/src/main/scala/cats/syntax/set.scala | 52 +++++++++++++++++++ .../scala/cats/tests/NonEmptySetSuite.scala | 5 ++ .../test/scala/cats/tests/SyntaxSuite.scala | 21 +++++++- 6 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 core/src/main/scala/cats/syntax/set.scala diff --git a/core/src/main/scala/cats/data/NonEmptySet.scala b/core/src/main/scala/cats/data/NonEmptySet.scala index ade9b6184ff..8ad0a4ed9a8 100644 --- a/core/src/main/scala/cats/data/NonEmptySet.scala +++ b/core/src/main/scala/cats/data/NonEmptySet.scala @@ -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 { diff --git a/core/src/main/scala/cats/syntax/all.scala b/core/src/main/scala/cats/syntax/all.scala index a8ce68e849e..971a53b55a1 100644 --- a/core/src/main/scala/cats/syntax/all.scala +++ b/core/src/main/scala/cats/syntax/all.scala @@ -64,3 +64,4 @@ trait AllSyntaxBinCompat1 with NestedSyntax with BinestedSyntax with ParallelFlatSyntax + with SetSyntax diff --git a/core/src/main/scala/cats/syntax/package.scala b/core/src/main/scala/cats/syntax/package.scala index bf95218ba10..2e925afae6c 100644 --- a/core/src/main/scala/cats/syntax/package.scala +++ b/core/src/main/scala/cats/syntax/package.scala @@ -52,4 +52,5 @@ package object syntax { object validated extends ValidatedSyntax object vector extends VectorSyntax object writer extends WriterSyntax + object set extends SetSyntax } diff --git a/core/src/main/scala/cats/syntax/set.scala b/core/src/main/scala/cats/syntax/set.scala new file mode 100644 index 00000000000..0fa5323537f --- /dev/null +++ b/core/src/main/scala/cats/syntax/set.scala @@ -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) + } +} diff --git a/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala b/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala index 36250049e4c..30b5dee8957 100644 --- a/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala +++ b/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala @@ -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)) + } + } } diff --git a/tests/src/test/scala/cats/tests/SyntaxSuite.scala b/tests/src/test/scala/cats/tests/SyntaxSuite.scala index c44b4d25051..8c392034c87 100644 --- a/tests/src/test/scala/cats/tests/SyntaxSuite.scala +++ b/tests/src/test/scala/cats/tests/SyntaxSuite.scala @@ -1,8 +1,10 @@ package cats package tests +import scala.collection.immutable.SortedSet +import scala.collection.immutable.SortedMap import cats.arrow.Compose -import cats.data.{Binested, Nested} +import cats.data.{Binested, Nested, NonEmptyList, NonEmptySet} import cats.instances.AllInstances import cats.syntax.{AllSyntax, AllSyntaxBinCompat} @@ -358,5 +360,22 @@ object SyntaxSuite extends AllSyntaxBinCompat with AllInstances with AllSyntax { val binested: Binested[F, G, H, A, B] = fgahb.binested } + + def testNonEmptySet[A: Order, B: Order] : Unit = { + val f = mock[A => B] + val set = mock[SortedSet[A]] + + val nes: Option[NonEmptySet[A]] = set.toNes + val grouped: SortedMap[B, NonEmptySet[A]] = set.groupByNes(f) + } + + def testNonEmptyList[A: Order, B: Order] : Unit = { + val f = mock[A => B] + val list = mock[List[A]] + + val nel: Option[NonEmptyList[A]] = list.toNel + val grouped: SortedMap[B, NonEmptyList[A]] = list.groupByNel(f) + } + }