Merge pull request #7 from palantir/stream_a_multimap

Make it possible to stream a Multimap
palantir · Jun 3, 2016 · 59145e2 · 59145e2
2 parents df79f11 + bfe69eb
commit 59145e2
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/main/java/com/palantir/common/streams/KeyedStream.java b/src/main/java/com/palantir/common/streams/KeyedStream.java
@@ -211,6 +211,13 @@ static <K, V> KeyedStream<K, V> stream(Map<K, V> map) {
         return new KeyedStreamImpl<>(map.entrySet().stream());
     }
 
+    /**
+     * Returns a keyed stream of {@code multimap}'s entries.
+     */
+    static <K, V> KeyedStream<K, V> stream(Multimap<K, V> multimap) {
+        return new KeyedStreamImpl<>(multimap.entries().stream());
+    }
+
     /**
      * Collects a stream and restreams it as a keyed stream, where key and value are the same.
      *

diff --git a/src/test/java/com/palantir/common/streams/KeyedStreamTests.java b/src/test/java/com/palantir/common/streams/KeyedStreamTests.java
@@ -31,8 +31,9 @@
 import org.junit.rules.ExpectedException;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Iterables;
-import com.palantir.common.streams.KeyedStream;
+import com.google.common.collect.SetMultimap;
 
 public class KeyedStreamTests {
 
@@ -69,6 +70,12 @@ public void test_stream_multiple_values_as_map() {
         assertThat(map).isEqualTo(ImmutableMap.of(3, 6, 4, 8, 6, 12));
     }
 
+    @Test
+    public void test_stream_multiple_values_as_multimap() {
+        SetMultimap<Integer, Integer> map = KeyedStream.stream(ImmutableSetMultimap.of(1, 2, 3, 4)).collectToSetMultimap();
+        assertThat(map).isEqualTo(ImmutableSetMultimap.of(3, 4, 1, 2));
+    }
+
     @Test
     public void test_collect_multiple_values_as_sorted_map() {
         Map<String, Integer> map = KeyedStream.of(Stream.of(4, 3, 5))