diff --git a/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java b/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java index f5aff4e15..318c829bc 100644 --- a/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java +++ b/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java @@ -548,6 +548,19 @@ public static BiCollector 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 BiCollector collectingAndThen( + BiCollector> collector, + BiFunction 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. diff --git a/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java b/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java index 3c77344e0..7d0f16924 100644 --- a/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java +++ b/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 Collector collectingAndThen( + Collector> upstream, + BiFunction 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 c : cases) { if (c.appliesTo(list)) { diff --git a/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java b/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java index 049d50125..0042f758c 100644 --- a/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java +++ b/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 salaries = BiStream.of("Joe", 1, "Tom", 2); assertThat(salaries.collect(inverse(toMap()))) diff --git a/mug/src/test/java/com/google/mu/util/stream/MoreCollectorsTest.java b/mug/src/test/java/com/google/mu/util/stream/MoreCollectorsTest.java index d3b9bbdd5..d2db1690b 100644 --- a/mug/src/test/java/com/google/mu/util/stream/MoreCollectorsTest.java +++ b/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);