Add updatedWith to Map

compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala

@@ -271,6 +271,10 @@ private[compat] trait PackageShared {
  self: scala.collection.Map[K, V]): MapExtensionMethods[K, V] =
  new MapExtensionMethods[K, V](self)
 
+ implicit def toImmutableMapExtensionMethods[K, V](
+ self: scala.collection.immutable.Map[K, V]): ImmutableMapExtensionMethods[K, V] =
+ new ImmutableMapExtensionMethods[K, V](self)
+
  implicit def toMutableMapExtensionMethods[K, V](
  self: scala.collection.mutable.Map[K, V]): MutableMapExtensionMethods[K, V] =
  new MutableMapExtensionMethods[K, V](self)
@@ -527,6 +531,16 @@ class MapExtensionMethods[K, V](private val self: scala.collection.Map[K, V]) ex
 
 }
 
+class ImmutableMapExtensionMethods[K, V](private val self: scala.collection.immutable.Map[K, V])
+ extends AnyVal {
+
+ def updatedWith[V1 >: V](key: K)(remappingFunction: (Option[V]) => Option[V1]): Map[K, V1] =
+ remappingFunction(self.get(key)) match {
+ case Some(v) => self.updated(key, v)
+ case None => self - key
+ }
+}
+
 class MutableMapExtensionMethods[K, V](private val self: scala.collection.mutable.Map[K, V])
  extends AnyVal {
 

compat/src/test/scala/test/scala/collection/MapTest.scala

@@ -16,6 +16,7 @@ import org.junit.Test
 import org.junit.Assert._
 
 import scala.collection.compat._
+import scala.collection.{immutable => i, mutable => m}
 
 class MapTest {
 
@@ -27,4 +28,43 @@ class MapTest {
  assertEquals(map, copy)
  }
 
+ @Test
+ def mutableUpdateWith: Unit = {
+ val map = m.Map("a" -> 1, "c" -> 3)
+ map.updateWith("b") {
+ case None => Some(2)
+ case Some(_) => Some(-1)
+ }
+ map.updateWith("c") {
+ case None => Some(-1)
+ case Some(_) => None
+ }
+ // unmodified entries are preserved
+ assertEquals(map.get("a"), Some(1))
+ // updateWith can add entries
+ assertEquals(map.get("b"), Some(2))
+ // updateWith can remove entries
+ assertFalse(map.contains("c"))
+ }
+
+ @Test
+ def immutableUpdatedWith: Unit = {
+ val map = i.Map("a" -> 1, "c" -> 3)
+ val bAdded = map.updatedWith("b") {
+ case None => Some(2)
+ case Some(_) => Some(-1)
+ }
+ val cRemoved = map.updatedWith("c") {
+ case None => Some(-1)
+ case Some(_) => None
+ }
+ // unmodified entries are preserved
+ assertEquals(map.get("a"), bAdded.get("a"))
+ assertEquals(map.get("a"), cRemoved.get("a"))
+ // updatedWith can add entries
+ assertEquals(bAdded.get("b"), Some(2))
+ // updatedWith can remove entries
+ assertFalse(cRemoved.contains("c"))
+ }
+
 }