From 6c05f3af7b306f426e7bf7006da0e6bf21123125 Mon Sep 17 00:00:00 2001
From: Samuel Souza <fsamuel-bs@users.noreply.github.com>
Date: Thu, 3 Nov 2022 17:38:31 +0000
Subject: [PATCH] Add peek methods (#127)

---
 changelog/@unreleased/pr-127.v2.yml                       | 5 +++++
 .../java/com/palantir/common/streams/KeyedStream.java     | 8 ++++++++
 .../java/com/palantir/common/streams/KeyedStreamImpl.java | 6 ++----
 3 files changed, 15 insertions(+), 4 deletions(-)
 create mode 100644 changelog/@unreleased/pr-127.v2.yml

diff --git a/changelog/@unreleased/pr-127.v2.yml b/changelog/@unreleased/pr-127.v2.yml
new file mode 100644
index 0000000..e0497db
--- /dev/null
+++ b/changelog/@unreleased/pr-127.v2.yml
@@ -0,0 +1,5 @@
+type: feature
+feature:
+  description: Add peek methods to KeyedStream
+  links:
+  - https://github.com/palantir/streams/pull/127
diff --git a/streams/src/main/java/com/palantir/common/streams/KeyedStream.java b/streams/src/main/java/com/palantir/common/streams/KeyedStream.java
index 889dc5c..0aed838 100644
--- a/streams/src/main/java/com/palantir/common/streams/KeyedStream.java
+++ b/streams/src/main/java/com/palantir/common/streams/KeyedStream.java
@@ -46,6 +46,14 @@
  */
 public interface KeyedStream<K, V> {
 
+    /**
+     * Returns a keyed stream consisting of elements in this stream,
+     * performing the provided action on each element of this stream.
+     */
+    default KeyedStream<K, V> peek(BiConsumer<K, V> consumer) {
+        return new KeyedStreamImpl<>(entries().peek(entry -> consumer.accept(entry.getKey(), entry.getValue())));
+    }
+
     /**
      * Returns a keyed stream consisting of the entries of this stream whose values match
      * the given predicate.
diff --git a/streams/src/main/java/com/palantir/common/streams/KeyedStreamImpl.java b/streams/src/main/java/com/palantir/common/streams/KeyedStreamImpl.java
index 2c2d4a8..8631be7 100644
--- a/streams/src/main/java/com/palantir/common/streams/KeyedStreamImpl.java
+++ b/streams/src/main/java/com/palantir/common/streams/KeyedStreamImpl.java
@@ -53,16 +53,14 @@ public <M extends Multimap<K, V>> M collectToMultimap(Supplier<M> supplier) {
     @Override
     public <K2, V2> KeyedStream<K2, V2> mapEntries(
             BiFunction<? super K, ? super V, ? extends Entry<? extends K2, ? extends V2>> entryMapper) {
-        return new KeyedStreamImpl<K2, V2>(entries.<Map.Entry<? extends K2, ? extends V2>>map(
-                entry -> entryMapper.apply(entry.getKey(), entry.getValue())));
+        return new KeyedStreamImpl<>(entries.map(entry -> entryMapper.apply(entry.getKey(), entry.getValue())));
     }
 
     @Override
     public <K2, V2> KeyedStream<K2, V2> flatMapEntries(
             BiFunction<? super K, ? super V, ? extends Stream<? extends Entry<? extends K2, ? extends V2>>>
                     entryMapper) {
-        return new KeyedStreamImpl<K2, V2>(
-                entries.flatMap(entry -> entryMapper.apply(entry.getKey(), entry.getValue())));
+        return new KeyedStreamImpl<>(entries.flatMap(entry -> entryMapper.apply(entry.getKey(), entry.getValue())));
     }
 
     @Override