Merge pull request #509 from emanb29/map-updatedwith

Add `updatedWith` to Map

Author: NthPortal

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"))
+  }
+
 }