collectingAndThen() for partitioningBy()

mug/src/main/java/com/google/mu/util/stream/BiCollectors.java

@@ -548,6 +548,19 @@ public static <K, V, T, R> BiCollector<K, V, R> collectingAndThen(
     };
   }
 
+  /**
+   * Returns a {@link BiCollector} that maps the result of {@code collector} using the {@code
+   * finisher} BiFunction. Useful when combined with BiCollectors like {@link #partitioningBy}.
+   *
+   * @since 8.1
+   */
+  public static <K, V, A, B, R> BiCollector<K, V, R> collectingAndThen(
+      BiCollector<K, V, ? extends Both<? extends A, ? extends B>> collector,
+      BiFunction<? super A, ? super B, ? extends R> finisher) {
+    requireNonNull(finisher);
+    return collectingAndThen(collector, ab -> ab.andThen(finisher));
+  }
+
   /**
    * Returns a {@link BiCollector} that first collects the input pairs into a {@link BiStream} and then applies
    * {@code finisher} on the intermediary BiStream.

mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java

@@ -16,7 +16,6 @@
 
 import static java.lang.Math.max;
 import static java.util.Objects.requireNonNull;
-import static java.util.stream.Collectors.collectingAndThen;
 import static java.util.stream.Collectors.toList;
 
 import java.util.ArrayList;
@@ -198,6 +197,19 @@ M build() {
     });
   }
 
+  /**
+   * Returns a {@link Collector} that maps the result of {@code upstream} collector using the {@code
+   * finisher} BiFunction. Useful when combined with collectors like {@link #partitioningBy}.
+   *
+   * @since 8.1
+   */
+  public static <T, A, B, R> Collector<T, ?, R> collectingAndThen(
+      Collector<T, ?, ? extends Both<? extends A, ? extends B>> upstream,
+      BiFunction<? super A, ? super B, ? extends R> finisher) {
+    requireNonNull(finisher);
+    return Collectors.collectingAndThen(upstream, ab -> ab.andThen(finisher));
+  }
+
   /**
    * Returns a collector that collects the only two elements from the input and transforms them
    * using the {@code mapper} function. If there are fewer or more elements in the input,
@@ -717,7 +729,7 @@ R build() {
     if (cases.size() == 1) {
       return cases.get(0);
     }
-    return collectingAndThen(toList(), list -> {
+    return Collectors.collectingAndThen(toList(), list -> {
       int elementsToShow = 1;
       for (FixedSizeCollector<T, ?, R> c : cases) {
         if (c.appliesTo(list)) {

mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java

@@ -380,6 +380,16 @@ public class BiCollectorsTest {
         .inOrder();
   }
 
+  @Test public void testCollectingAndThen_fomPair() {
+    String result =
+        BiStream.of(1, "one", 2, "two", 3, "three", 4, "four", 5, "five")
+            .collect(
+                collectingAndThen(
+                    partitioningBy((i, n) -> i % 2 == 1),
+                    (odds, evens) -> "odd:" + odds.toMap() + "; even:" + evens.toMap()));
+    assertThat(result).isEqualTo("odd:{1=one, 3=three, 5=five}; even:{2=two, 4=four}");
+  }
+
   @Test public void testInverse_toStream() {
     BiStream<String, Integer> salaries = BiStream.of("Joe", 1, "Tom", 2);
     assertThat(salaries.collect(inverse(toMap())))

mug/src/test/java/com/google/mu/util/stream/MoreCollectorsTest.java

@@ -24,6 +24,7 @@
 import static com.google.mu.util.stream.BiCollectors.toMap;
 import static com.google.mu.util.stream.MoreCollectors.allMax;
 import static com.google.mu.util.stream.MoreCollectors.allMin;
+import static com.google.mu.util.stream.MoreCollectors.collectingAndThen;
 import static com.google.mu.util.stream.MoreCollectors.combining;
 import static com.google.mu.util.stream.MoreCollectors.flatMapping;
 import static com.google.mu.util.stream.MoreCollectors.flatteningMaps;
@@ -155,6 +156,17 @@ public class MoreCollectorsTest {
     assertThrows(UnsupportedOperationException.class, list::clear);
   }
 
+  @Test public void testCollectingAndThen_fromPair() {
+    String result =
+        Stream.of(1, 2, 3, 4, 5)
+            .collect(
+                collectingAndThen(
+                    partitioningBy(
+                        n -> n % 2 == 1, toImmutableList(), summingInt(Integer::intValue)),
+                    (odd, even) -> "odd:" + odd + "; sum of even:" + even));
+    assertThat(result).isEqualTo("odd:[1, 3, 5]; sum of even:6");
+  }
+
   @Test
   public void testCombining_two() {
     assertThat(Stream.of(1, 10).collect(combining(Integer::sum)).intValue()).isEqualTo(11);