diff --git a/core/src/main/scala/cats/data/NonEmptyMapImpl.scala b/core/src/main/scala/cats/data/NonEmptyMapImpl.scala index 8b8479cb90..3510da6be8 100644 --- a/core/src/main/scala/cats/data/NonEmptyMapImpl.scala +++ b/core/src/main/scala/cats/data/NonEmptyMapImpl.scala @@ -93,6 +93,15 @@ sealed class NonEmptyMapOps[K, A](val value: NonEmptyMap[K, A]) { */ def lookup(k: K): Option[A] = toSortedMap.get(k) + /** + * Applies f to the value stored at k. If lookup misses, does nothing. + */ + def updateWith(k: K)(f: A => A): NonEmptyMap[K, A] = + lookup(k) match { + case Some(v) => add((k, f(v))) + case None => value + } + /** * Returns a `SortedSet` containing all the keys of this map. */ diff --git a/tests/src/test/scala/cats/tests/NonEmptyMapSuite.scala b/tests/src/test/scala/cats/tests/NonEmptyMapSuite.scala index 9b11b522af..6b0af5cad6 100644 --- a/tests/src/test/scala/cats/tests/NonEmptyMapSuite.scala +++ b/tests/src/test/scala/cats/tests/NonEmptyMapSuite.scala @@ -209,4 +209,22 @@ class NonEmptyMapSuite extends CatsSuite { nem.toNel should ===(NonEmptyList.fromListUnsafe(nem.toSortedMap.toList)) } } + + test("NonEmptyMap#updateWith identity should be a no-op") { + forAll { (nem: NonEmptyMap[String, Int], i: (String, Int)) => + nem.add(i) should ===(nem.add(i).updateWith(i._1)(identity)) + } + } + + test("NonEmptyMap#updateWith on existing value should behave like Option#map on the same value") { + forAll { (nem: NonEmptyMap[String, Int], i: (String, Int)) => + nem.add(i).lookup(i._1).map(_ + 1) should ===(nem.add(i).updateWith(i._1)(_ + 1).lookup(i._1)) + } + } + + test("NonEmptyMap#updateWith should not act when key is missing") { + val single = NonEmptyMap[String, Int](("here", 1), SortedMap()) + single.lookup("notHere") should ===(single.updateWith("notHere")(_ => 1).lookup("notHere")) + } + }