feat(ast): add support for LambdaExpr to infer type from return expr type (#1011)

zhumin8 · web-flow · commit a1795582948d · 2022-07-11T13:46:00.000-04:00
This will enable it to be assigned to an appropriate variable

Before this change: cannot assign a lambda expression to a variable because its type is always VOID.

After this change: lambda expression builder will infer type from returnExpr’s expression type. Allowing it to be assigned to an appropriate variable. Because returnExpr is a required field, its type should always be accessible.

Example: a naive sample usage is provided in JavaWriterVisitorTest.java as test.
diff --git a/src/main/java/com/google/api/generator/engine/ast/LambdaExpr.java b/src/main/java/com/google/api/generator/engine/ast/LambdaExpr.java
@@ -27,9 +27,7 @@
 public abstract class LambdaExpr implements Expr {
   @Override
   public TypeNode type() {
-    // TODO(v2): Support set of FunctionalInterface  parameterized on the args and return type,
-    // which would enable assignment to an appropriate variable.
-    return TypeNode.VOID;
+    return returnExpr().expr().type();
   }
 
   public abstract ImmutableList<VariableExpr> arguments();
diff --git a/src/test/java/com/google/api/generator/engine/ast/LambdaExprTest.java b/src/test/java/com/google/api/generator/engine/ast/LambdaExprTest.java
@@ -14,6 +14,7 @@
 
 package com.google.api.generator.engine.ast;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 
 import java.util.Arrays;
@@ -27,6 +28,15 @@ public void validLambdaExpr_noArguments() {
         .build();
   }
 
+  @Test
+  public void validLambdaExpr_inferTypeFromReturnExpr() {
+    LambdaExpr lambdaExpr =
+        LambdaExpr.builder()
+            .setReturnExpr(ValueExpr.withValue(StringObjectValue.withValue("foo")))
+            .build();
+    assertEquals(TypeNode.STRING, lambdaExpr.type());
+  }
+
   @Test
   public void validLambdaExpr_severalArguments() {
     VariableExpr argOneVarExpr =
diff --git a/src/test/java/com/google/api/generator/engine/writer/JavaWriterVisitorTest.java b/src/test/java/com/google/api/generator/engine/writer/JavaWriterVisitorTest.java
@@ -2345,6 +2345,23 @@ public void writeLambdaExpr_noParameters() {
     assertEquals("() -> \"foo\"", writerVisitor.write());
   }
 
+  @Test
+  public void writeLambdaExpr_assignToVariable() {
+    LambdaExpr lambdaExpr =
+        LambdaExpr.builder()
+            .setReturnExpr(ValueExpr.withValue(StringObjectValue.withValue("foo")))
+            .build();
+    AssignmentExpr assignmentExpr =
+        AssignmentExpr.builder()
+            .setVariableExpr(
+                VariableExpr.withVariable(
+                    Variable.builder().setName("word").setType(TypeNode.STRING).build()))
+            .setValueExpr(lambdaExpr)
+            .build();
+    assignmentExpr.accept(writerVisitor);
+    assertEquals("word = () -> \"foo\"", writerVisitor.write());
+  }
+
   @Test
   public void writeLambdaExpr_oneParameter() {
     VariableExpr argVarExpr =
