fix CollectHeadOption

xuwei-k · Dec 7, 2024 · 3dd9cfe · 3dd9cfe
1 parent ed5f351
commit 3dd9cfe
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/input/src/main/scala/fix/CollectHeadOptionTest.scala b/input/src/main/scala/fix/CollectHeadOptionTest.scala
@@ -5,4 +5,5 @@ package fix
 
 object CollectHeadOptionTest {
   def x1: Option[String] = List(1, 2, 3).collect { case n if n % 2 == 0 => n.toString }.headOption
+  def x2(f: PartialFunction[Int, String]): Option[String] = List(1, 2, 3).collect(f).headOption
 }
diff --git a/output/src/main/scala/fix/CollectHeadOptionTest.scala b/output/src/main/scala/fix/CollectHeadOptionTest.scala
@@ -2,4 +2,5 @@ package fix
 
 object CollectHeadOptionTest {
   def x1: Option[String] = List(1, 2, 3).collectFirst{ case n if n % 2 == 0 => n.toString }
+  def x2(f: PartialFunction[Int, String]): Option[String] = List(1, 2, 3).collectFirst(f)
 }
diff --git a/rules/src/main/scala/fix/CollectHeadOption.scala b/rules/src/main/scala/fix/CollectHeadOption.scala
@@ -17,7 +17,12 @@ class CollectHeadOption extends SyntacticRule("CollectHeadOption") {
             ),
             Term.Name("headOption")
           ) =>
-        Patch.replaceTree(t, s"${obj}.collectFirst${func}")
+        func match {
+          case _: Term.PartialFunction =>
+            Patch.replaceTree(t, s"${obj}.collectFirst${func}")
+          case _ =>
+            Patch.replaceTree(t, s"${obj}.collectFirst(${func})")
+        }
     }
   }.asPatch
 }