From 01c7a2f18ab6dbdd170dc29f63a6fcc03334a0f3 Mon Sep 17 00:00:00 2001
From: ncreep <daniel.beskin@gmail.com>
Date: Tue, 6 Sep 2022 12:42:15 +0300
Subject: [PATCH] Adding a `MapPartial` wart

---
 .../org/wartremover/warts/MapPartial.scala    | 25 +++++++++
 .../org/wartremover/test/MapPartialTest.scala | 56 +++++++++++++++++++
 docs/_posts/2017-02-11-warts.md               |  6 ++
 3 files changed, 87 insertions(+)
 create mode 100644 core/src/main/scala-2/org/wartremover/warts/MapPartial.scala
 create mode 100644 core/src/test/scala/org/wartremover/test/MapPartialTest.scala

diff --git a/core/src/main/scala-2/org/wartremover/warts/MapPartial.scala b/core/src/main/scala-2/org/wartremover/warts/MapPartial.scala
new file mode 100644
index 00000000..75bc33b6
--- /dev/null
+++ b/core/src/main/scala-2/org/wartremover/warts/MapPartial.scala
@@ -0,0 +1,25 @@
+package org.wartremover.warts
+
+import org.wartremover.{WartTraverser, WartUniverse}
+
+object MapPartial extends WartTraverser {
+  def apply(u: WartUniverse): u.Traverser = {
+    import u.universe._
+
+    val mapSymbol = rootMirror.staticClass("scala.collection.Map")
+    val ApplyName = TermName("apply")
+    new u.Traverser {
+      override def traverse(tree: Tree): Unit = {
+        tree match {
+          // Ignore trees marked by SuppressWarnings
+          case t if hasWartAnnotation(u)(t) =>
+          case Select(left, ApplyName) if left.tpe.baseType(mapSymbol) != NoType =>
+            error(u)(tree.pos, "Map#apply is disabled - use Map#getOrElse instead")
+          case LabelDef(_, _, rhs) if isSynthetic(u)(tree) =>
+          case _ =>
+            super.traverse(tree)
+        }
+      }
+    }
+  }
+}
diff --git a/core/src/test/scala/org/wartremover/test/MapPartialTest.scala b/core/src/test/scala/org/wartremover/test/MapPartialTest.scala
new file mode 100644
index 00000000..d91c72ef
--- /dev/null
+++ b/core/src/test/scala/org/wartremover/test/MapPartialTest.scala
@@ -0,0 +1,56 @@
+package org.wartremover.test
+
+import org.scalatest.Inspectors
+import org.scalatest.funsuite.AnyFunSuite
+import org.wartremover.warts.MapPartial
+
+class MapPartialTest extends AnyFunSuite with ResultAssertions with Inspectors {
+  test("can't use Map#apply on all kinds of Maps") {
+    val result1 = WartTestTraverser(MapPartial) {
+      println(Map("a" -> 1)("a"))
+    }
+    val result2 = WartTestTraverser(MapPartial) {
+      println(collection.Map("a" -> 1)("a"))
+    }
+    val result3 = WartTestTraverser(MapPartial) {
+      println(collection.immutable.Map("a" -> 1)("a"))
+    }
+    val result4 = WartTestTraverser(MapPartial) {
+      println(collection.mutable.Map("a" -> 1)("a"))
+    }
+
+    forAll(List(result1, result2, result3, result4)) { result =>
+      assertError(result)("Map#apply is disabled - use Map#getOrElse instead")
+    }
+  }
+
+  test("can't use Map#apply without syntax-sugar") {
+    val result = WartTestTraverser(MapPartial) {
+      println(Map("a" -> 1).apply("a"))
+    }
+
+    assertError(result)("Map#apply is disabled - use Map#getOrElse instead")
+  }
+
+  test("doesn't detect other `apply` methods") {
+    val result = WartTestTraverser(MapPartial) {
+      class A {
+        def apply(i: Int) = 3
+      }
+      println((new A).apply(2))
+    }
+
+    assertEmpty(result)
+  }
+
+  test("MapPartial wart obeys SuppressWarnings") {
+    val result = WartTestTraverser(MapPartial) {
+      @SuppressWarnings(Array("org.wartremover.warts.MapPartial"))
+      val foo = {
+        println(Map("a" -> 1)("a"))
+      }
+    }
+
+    assertEmpty(result)
+  }
+}
diff --git a/docs/_posts/2017-02-11-warts.md b/docs/_posts/2017-02-11-warts.md
index 7ae117ad..d1d99089 100644
--- a/docs/_posts/2017-02-11-warts.md
+++ b/docs/_posts/2017-02-11-warts.md
@@ -178,6 +178,12 @@ file 2:
 class d extends c
 ```
 
+### MapPartial
+
+`scala.collection.Map` has an `apply` method which will throw if the key is missing.
+The program should be refactored to use `scala.collection.Map#getOrElse` to
+explicitly handle the case of a missing key.
+
 ### MutableDataStructures
 
 The standard library provides mutable collections. Mutation breaks equational reasoning.