Refactor: Parse return statements without an expression properly

This is an internal refactoring necessary for the Skylark linter.
It does not change any behavior.

RELNOTES: None
PiperOrigin-RevId: 166199367

2 files changed, 6 insertions, 12 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/syntax/ASTPrettyPrintTest.java b/src/test/java/com/google/devtools/build/lib/syntax/ASTPrettyPrintTest.java
index 7c7b02acfa..8278345aa3 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/ASTPrettyPrintTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/ASTPrettyPrintTest.java
@@ -393,12 +393,11 @@ public class ASTPrettyPrintTest extends EvaluationTestCase {
         new ReturnStatement(new StringLiteral("foo")),
         "return \"foo\"\n");
 
-    assertIndentedPrettyMatches(
-        new ReturnStatement(new Identifier("None")),
-        "  return\n");
-    assertTostringMatches(
-        new ReturnStatement(new Identifier("None")),
-        "return\n");
+    assertIndentedPrettyMatches(new ReturnStatement(new Identifier("None")), "  return None\n");
+    assertTostringMatches(new ReturnStatement(new Identifier("None")), "return None\n");
+
+    assertIndentedPrettyMatches(new ReturnStatement(null), "  return\n");
+    assertTostringMatches(new ReturnStatement(null), "return\n");
   }
 
   // Miscellaneous.
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
index 94cdb30105..288038f4ef 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
@@ -971,12 +971,7 @@ public class ParserTest extends EvaluationTestCase {
       assertThat(bodyNoExpr).hasSize(1);
 
       ReturnStatement returnNoExpr = (ReturnStatement) bodyNoExpr.get(0);
-      Identifier none = (Identifier) returnNoExpr.getReturnExpression();
-      assertThat(none.getName()).isEqualTo("None");
-      assertLocation(
-          returnNoExpr.getLocation().getStartOffset(),
-          returnNoExpr.getLocation().getEndOffset(),
-          none.getLocation());
+      assertThat(returnNoExpr.getReturnExpression()).isNull();
     }
   }