Support replacing guava's newArrayListWithCapacity

openrewrite · Jun 16, 2021 · d4be80f · d4be80f
1 parent 26b81d2
commit d4be80f
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/src/main/java/org/openrewrite/java/migrate/guava/NoGuavaListsNewArrayList.java b/src/main/java/org/openrewrite/java/migrate/guava/NoGuavaListsNewArrayList.java
@@ -28,6 +28,7 @@
 public class NoGuavaListsNewArrayList extends Recipe {
     private static final MethodMatcher NEW_ARRAY_LIST = new MethodMatcher("com.google.common.collect.Lists newArrayList()");
     private static final MethodMatcher NEW_ARRAY_LIST_ITERABLE = new MethodMatcher("com.google.common.collect.Lists newArrayList(java.lang.Iterable)");
+    private static final MethodMatcher NEW_ARRAY_LIST_CAPACITY = new MethodMatcher("com.google.common.collect.Lists newArrayListWithCapacity(int)");
 
     @Override
     public String getDisplayName() {
@@ -46,6 +47,7 @@ protected TreeVisitor<?, ExecutionContext> getApplicableTest() {
             public J.CompilationUnit visitCompilationUnit(J.CompilationUnit cu, ExecutionContext executionContext) {
                 doAfterVisit(new UsesMethod<>(NEW_ARRAY_LIST));
                 doAfterVisit(new UsesMethod<>(NEW_ARRAY_LIST_ITERABLE));
+                doAfterVisit(new UsesMethod<>(NEW_ARRAY_LIST_CAPACITY));
                 return cu;
             }
         };
@@ -62,6 +64,10 @@ protected TreeVisitor<?, ExecutionContext> getVisitor() {
                     .imports("java.util.ArrayList")
                     .build();
 
+            private final JavaTemplate newArrayListCapacity = template("new ArrayList<>(#{any(int)})")
+                    .imports("java.util.ArrayList")
+                    .build();
+
             @Override
             public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext executionContext) {
                 if (NEW_ARRAY_LIST.matches(method)) {
@@ -73,6 +79,11 @@ public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext execu
                     maybeAddImport("java.util.ArrayList");
                     return method.withTemplate(newArrayListIterable, method.getCoordinates().replace(),
                             method.getArguments().get(0));
+                } else if (NEW_ARRAY_LIST_CAPACITY.matches(method)) {
+                    maybeRemoveImport("com.google.common.collect.Lists");
+                    maybeAddImport("java.util.ArrayList");
+                    return method.withTemplate(newArrayListCapacity, method.getCoordinates().replace(),
+                            method.getArguments().get(0));
                 }
                 return super.visitMethodInvocation(method, executionContext);
             }

diff --git a/src/test/kotlin/org/openrewrite/java/migrate/guava/NoGuavaListsNewArrayListTest.kt b/src/test/kotlin/org/openrewrite/java/migrate/guava/NoGuavaListsNewArrayListTest.kt
@@ -76,4 +76,26 @@ class NoGuavaListsNewArrayListTest: JavaRecipeTest {
             }
         """
     )
+
+    @Test
+    fun replaceWithNewArrayListWithCapacity() = assertChanged(
+        before = """
+            import com.google.common.collect.*;
+            
+            import java.util.ArrayList;
+            import java.util.List;
+            
+            class Test {
+                List<Integer> cardinalsWorldSeries = Lists.newArrayListWithCapacity(2);
+            }
+        """,
+        after = """
+            import java.util.ArrayList;
+            import java.util.List;
+            
+            class Test {
+                List<Integer> cardinalsWorldSeries = new ArrayList<>(2);
+            }
+        """
+    )
 }