add groupByNeSeq and groupByNeSeqA to SeqOps

typelevel · Sep 25, 2021 · c70efc8 · c70efc8
1 parent e70f22a
commit c70efc8
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 16 deletions.
diff --git a/core/src/main/scala/cats/syntax/seq.scala b/core/src/main/scala/cats/syntax/seq.scala
@@ -1,7 +1,13 @@
 package cats.syntax
 
+import cats.Applicative
+import cats.Functor
+import cats.Order
+import cats.Traverse
 import cats.data.NonEmptySeq
+
 import scala.collection.immutable.Seq
+import scala.collection.immutable.SortedMap
 
 trait SeqSyntax {
   implicit final def catsSyntaxSeqs[A](va: Seq[A]): SeqOps[A] = new SeqOps(va)
@@ -43,4 +49,61 @@ final class SeqOps[A](private val va: Seq[A]) extends AnyVal {
    * }}}
    */
   def concatNeSeq[AA >: A](neseq: NonEmptySeq[AA]): NonEmptySeq[AA] = neseq.prependSeq(va)
+
+  /**
+   * Groups elements inside this `Seq` according to the `Order` of the keys
+   * produced by the given mapping function.
+   *
+   * {{{
+   * scala> import cats.data.NonEmptySeq
+   * scala> import cats.syntax.all._
+   * scala> import scala.collection.immutable.Seq
+   * scala> import scala.collection.immutable.SortedMap
+   *
+   * scala> val seq = Seq(12, -2, 3, -5)
+   * scala> val res = SortedMap(false -> NonEmptySeq.of(-2, -5), true -> NonEmptySeq.of(12, 3))
+   * scala> seq.groupByNeSeq(_ >= 0) === res
+   * res0: Boolean = true
+   * }}}
+   */
+  def groupByNeSeq[B](f: A => B)(implicit B: Order[B]): SortedMap[B, NonEmptySeq[A]] = {
+    implicit val ordering: Ordering[B] = B.toOrdering
+    toNeSeq.fold(SortedMap.empty[B, NonEmptySeq[A]])(_.groupBy(f))
+  }
+
+  /**
+   * Groups elements inside this `Seq` according to the `Order` of the keys
+   * produced by the given mapping monadic function.
+   *
+   * {{{
+   * scala> import cats.data.NonEmptySeq
+   * scala> import cats.syntax.all._
+   * scala> import scala.collection.immutable.Seq
+   * scala> import scala.collection.immutable.SortedMap
+   *
+   * scala> def f(n: Int) = n match { case 0 => None; case n => Some(n > 0) }
+   *
+   * scala> val seq = Seq(12, -2, 3, -5)
+   * scala> val res = Some(SortedMap(false -> NonEmptySeq.of(-2, -5), true -> NonEmptySeq.of(12, 3)))
+   * scala> seq.groupByNeSeqA(f) === res
+   * res0: Boolean = true
+   *
+   * scala> // `f(0)` returns `None`
+   * scala> (seq :+ 0).groupByNeSeqA(f) === None
+   * res1: Boolean = true
+   * }}}
+   */
+  def groupByNeSeqA[F[_], B](
+    f: A => F[B]
+  )(implicit F: Applicative[F], B: Order[B]): F[SortedMap[B, NonEmptySeq[A]]] = {
+    implicit val ordering: Ordering[B] = B.toOrdering
+    val mapFunctor = Functor[SortedMap[B, *]]
+    val nesTraverse = Traverse[NonEmptySeq]
+
+    toNeSeq.fold(F.pure(SortedMap.empty[B, NonEmptySeq[A]])) { nes =>
+      F.map(nesTraverse.traverse(nes)(a => F.tupleRight(f(a), a))) { seq =>
+        mapFunctor.map(seq.groupBy(_._1))(_.map(_._2))
+      }
+    }
+  }
 }
diff --git a/tests/src/test/scala/cats/tests/SeqSuite.scala b/tests/src/test/scala/cats/tests/SeqSuite.scala
@@ -1,24 +1,18 @@
 package cats.tests
 
-import cats.{Align, Alternative, CoflatMap, Monad, Semigroupal, Traverse, TraverseFilter}
-import cats.data.{NonEmptySeq, ZipSeq}
-import cats.laws.discipline.{
-  AlignTests,
-  AlternativeTests,
-  CoflatMapTests,
-  CommutativeApplyTests,
-  MonadTests,
-  SemigroupalTests,
-  SerializableTests,
-  ShortCircuitingTests,
-  TraverseFilterTests,
-  TraverseTests
-}
+import cats._
+import cats.data.NonEmptySeq
+import cats.data.OneAnd
+import cats.data.ZipSeq
+import cats.laws.discipline._
 import cats.laws.discipline.arbitrary._
-import cats.syntax.show._
-import cats.syntax.seq._
 import cats.syntax.eq._
+import cats.syntax.functor._
+import cats.syntax.seq._
+import cats.syntax.show._
+import cats.syntax.traverse._
 import org.scalacheck.Prop._
+
 import scala.collection.immutable.Seq
 
 class SeqSuite extends CatsSuite {
@@ -71,6 +65,36 @@ class SeqSuite extends CatsSuite {
     }
   }
 
+  test("groupByNeSeq should be consistent with groupBy") {
+    forAll { (fa: Seq[String], f: String => Int) =>
+      assert((fa.groupByNeSeq(f).map { _.map(_.toSeq) }: Map[Int, Seq[String]]) === fa.groupBy(f))
+    }
+  }
+
+  test("groupByNeSeqA[Id] should be consistent with groupByNeSeq") {
+    forAll { (fa: Seq[String], f: String => Id[Int]) =>
+      assert(fa.groupByNeSeqA(f) === fa.groupByNeSeq(f))
+    }
+  }
+
+  def unwrapGroupByNeSeq[G[_]: Applicative, A, B](fa: Seq[A], f: A => G[B])(implicit ord: Order[B]) = {
+    implicit val ordering = ord.toOrdering
+    fa
+      .traverse(a => f(a).tupleRight(a))
+      .map(_.groupByNeSeq(_._1).map(_.map(_.map(_._2))))
+  }
+
+  test("groupByNeSeqA[Option] should be consistent with groupByNeSeq") {
+    forAll { (fa: Seq[String], f: String => Option[Int]) =>
+      assert(fa.groupByNeSeqA(f) === unwrapGroupByNeSeq(fa, f))
+    }
+  }
+  test("groupByNeSeqA[OneAnd] should be consistent with groupByNeSeq") {
+    forAll { (fa: Seq[String], f: String => OneAnd[Option, Int]) =>
+      assert(fa.groupByNeSeqA(f) === unwrapGroupByNeSeq(fa, f))
+    }
+  }
+
   test("traverse is stack-safe") {
     val seq = (0 until 100000).toSeq
     val sumAll = Traverse[Seq]