Fix RemoveSystemOutPrintln for lambdas

openrewrite · Feb 20, 2025 · c54f423 · c54f423
1 parent 61f8b01
commit c54f423
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 6 deletions.
diff --git a/src/main/java/org/openrewrite/staticanalysis/RemoveSystemOutPrintln.java b/src/main/java/org/openrewrite/staticanalysis/RemoveSystemOutPrintln.java
@@ -16,14 +16,18 @@
 package org.openrewrite.staticanalysis;
 
 import org.jspecify.annotations.Nullable;
-import org.openrewrite.ExecutionContext;
-import org.openrewrite.Preconditions;
-import org.openrewrite.Recipe;
-import org.openrewrite.TreeVisitor;
+import org.openrewrite.*;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.MethodMatcher;
 import org.openrewrite.java.search.UsesMethod;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JRightPadded;
+import org.openrewrite.java.tree.Space;
+import org.openrewrite.marker.Markers;
+
+import java.util.Collections;
+
+import static org.openrewrite.Tree.randomId;
 
 public class RemoveSystemOutPrintln extends Recipe {
     private static final MethodMatcher SYSTEM_OUT_PRINTLN = new MethodMatcher("java.io.PrintStream println(..)");
@@ -43,9 +47,19 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
         return Preconditions.check(new UsesMethod<>(SYSTEM_OUT_PRINTLN), new JavaIsoVisitor<ExecutionContext>() {
 
             @Override
+            public J.Lambda visitLambda(J.Lambda lambda, ExecutionContext ctx) {
+                J.Lambda l = super.visitLambda(lambda, ctx);
+                //noinspection ConstantValue
+                if (l.getBody() == null) {
+                    l = l.withBody(new J.Block(randomId(), lambda.getPrefix(), Markers.EMPTY, JRightPadded.build(false), Collections.emptyList(), Space.EMPTY));
+                }
+                return l;
+            }
+
+            @Override
+            @SuppressWarnings("NullableProblems")
             public  J.@Nullable MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
                 if (SYSTEM_OUT_PRINTLN.matches(method)) {
-                    //noinspection DataFlowIssue
                     return null;
                 }
                 return super.visitMethodInvocation(method, ctx);

diff --git a/src/test/java/org/openrewrite/staticanalysis/RemoveSystemOutPrintlnTest.java b/src/test/java/org/openrewrite/staticanalysis/RemoveSystemOutPrintlnTest.java
@@ -38,7 +38,7 @@ void removePrintln() {
             """
               class Test {
                   void test() {
-                        System.out.println("Hello, world!");
+                      System.out.println("Hello, world!");
                   }
               }
               """,
@@ -51,4 +51,31 @@ void test() {
           )
         );
     }
+
+    @Test
+    void lambda() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.function.Consumer;
+
+              class Test {
+                  void test() {
+                      Consumer<String> c = s -> System.out.println(s);
+                  }
+              }
+              """,
+            """
+              import java.util.function.Consumer;
+
+              class Test {
+                  void test() {
+                      Consumer<String> c = s -> {};
+                  }
+              }
+              """
+          )
+        );
+    }
 }