Optimize AnyToNullable

openrewrite · Apr 30, 2024 · cf80960 · cf80960
1 parent 34d0dcb
commit cf80960
Showing 1 changed file with 10 additions and 11 deletions.
diff --git a/src/main/java/org/openrewrite/java/testing/mockito/AnyToNullable.java b/src/main/java/org/openrewrite/java/testing/mockito/AnyToNullable.java
@@ -20,6 +20,7 @@
 import org.openrewrite.java.ChangeMethodName;
 import org.openrewrite.java.ChangeMethodTargetToStatic;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaSourceFile;
 import org.openrewrite.xml.tree.Xml;
 
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -63,20 +64,18 @@ public TreeVisitor<?, ExecutionContext> getScanner(AtomicBoolean acc) {
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor(AtomicBoolean acc) {
-        ChangeMethodName changeMethodName = new ChangeMethodName(
-                "org.mockito.Mockito any(java.lang.Class)", "nullable", null, null);
-        ChangeMethodTargetToStatic changeMethodTargetToStatic = new ChangeMethodTargetToStatic(
-                "org.mockito.Mockito nullable(java.lang.Class)", "org.mockito.ArgumentMatchers", null, null);
-        AnyStringToNullable anyStringToNullable = new AnyStringToNullable();
         return Preconditions.check(acc.get(), new TreeVisitor<Tree, ExecutionContext>() {
             @Override
-            public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
-                if (tree instanceof J) {
-                    doAfterVisit(changeMethodName.getVisitor());
-                    doAfterVisit(changeMethodTargetToStatic.getVisitor());
-                    doAfterVisit(anyStringToNullable.getVisitor());
+            public @Nullable Tree preVisit(Tree tree, ExecutionContext ctx) {
+                if (tree instanceof JavaSourceFile) {
+                    stopAfterPreVisit();
+                    doAfterVisit(new ChangeMethodName(
+                            "org.mockito.Mockito any(java.lang.Class)", "nullable", null, null).getVisitor());
+                    doAfterVisit(new ChangeMethodTargetToStatic(
+                            "org.mockito.Mockito nullable(java.lang.Class)", "org.mockito.ArgumentMatchers", null, null).getVisitor());
+                    doAfterVisit(new AnyStringToNullable().getVisitor());
                 }
-                return super.visit(tree, ctx);
+                return super.preVisit(tree, ctx);
             }
         });
     }