fix

uber · Aug 25, 2023 · b3e3056 · b3e3056
1 parent e16779a
commit b3e3056
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java b/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
@@ -199,7 +199,7 @@ private Description.Builder addSuggestedSuppression(
       case ASSIGN_FIELD_NULLABLE:
       case SWITCH_EXPRESSION_NULLABLE:
         if (config.getCastToNonNullMethod() != null) {
-          builder = addCastToNonNullFix(suggestTree, builder);
+          builder = addCastToNonNullFix(suggestTree, builder, state);
         } else {
           builder = addSuppressWarningsFix(suggestTree, builder, suppressionName);
         }
@@ -320,15 +320,16 @@ private Tree suppressibleNode(@Nullable TreePath path) {
         .orElse(null);
   }
 
-  private Description.Builder addCastToNonNullFix(Tree suggestTree, Description.Builder builder) {
+  private Description.Builder addCastToNonNullFix(
+      Tree suggestTree, Description.Builder builder, VisitorState state) {
     final String fullMethodName = config.getCastToNonNullMethod();
     if (fullMethodName == null) {
       throw new IllegalStateException("cast-to-non-null method not set");
     }
     // Add a call to castToNonNull around suggestTree:
     final String[] parts = fullMethodName.split("\\.");
     final String shortMethodName = parts[parts.length - 1];
-    final String replacement = shortMethodName + "(" + suggestTree.toString() + ")";
+    final String replacement = shortMethodName + "(" + state.getSourceForNode(suggestTree) + ")";
     final SuggestedFix fix =
         SuggestedFix.builder()
             .replace(suggestTree, replacement)
@@ -354,7 +355,7 @@ private Description.Builder removeCastToNonNullFix(
             invTree, suggestTree));
     // Remove the call to castToNonNull:
     final SuggestedFix fix =
-        SuggestedFix.builder().replace(invTree, suggestTree.toString()).build();
+        SuggestedFix.builder().replace(invTree, state.getSourceForNode(suggestTree)).build();
     return builder.addFix(fix);
   }
 

diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java
@@ -22,6 +22,8 @@
 
 package com.uber.nullaway;
 
+import static com.google.errorprone.BugCheckerRefactoringTestHelper.TestMode.TEXT_MATCH;
+
 import com.google.errorprone.BugCheckerRefactoringTestHelper;
 import com.google.errorprone.ErrorProneFlags;
 import com.uber.nullaway.testlibrarymodels.TestLibraryModels;
@@ -161,6 +163,38 @@ public void removeUnnecessaryCastToNonNullFromLibraryModel() throws IOException
         .doTest();
   }
 
+  @Test
+  public void removeUnnecessaryCastToNonNullMultiLine() throws IOException {
+    makeTestHelper()
+        .addInputLines(
+            "Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "class Test {",
+            "  static class Foo { Object getObj() { return new Object(); } }",
+            "  Object test1(Foo f) {",
+            "    return castToNonNull(f",
+            "                         // comment that should not be deleted",
+            "                         .getObj());",
+            "  }",
+            "}")
+        .addOutputLines(
+            "out/Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "class Test {",
+            "  static class Foo { Object getObj() { return new Object(); } }",
+            "  Object test1(Foo f) {",
+            "    return f",
+            "        // comment that should not be deleted",
+            "        .getObj();",
+            "  }",
+            "}")
+        .doTest(TEXT_MATCH);
+  }
+
   @Test
   public void suggestSuppressionOnMethodRef() throws IOException {
     makeTestHelper()
@@ -222,6 +256,7 @@ public void suggestCastToNonNullMultiLine() throws IOException {
             "out/Test.java",
             "package com.uber;",
             "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "",
             "import javax.annotation.Nullable;",
             "class Test {",
             "  static class Foo { @Nullable Object getObj() { return null; } }",
@@ -231,6 +266,6 @@ public void suggestCastToNonNullMultiLine() throws IOException {
             "           .getObj());",
             "  }",
             "}")
-        .doTest();
+        .doTest(TEXT_MATCH);
   }
 }