Add Lambda support through Lambdaspec class

src/main/java/com/squareup/javapoet/CodeBlock.java

@@ -415,6 +415,11 @@ public Builder add(CodeBlock codeBlock) {
       return this;
     }
 
+    public Builder addLambda(LambdaSpec lambda) {
+      add(lambda.toString());
+      return this;
+    }
+
     public Builder indent() {
       this.formatParts.add("$>");
       return this;

src/main/java/com/squareup/javapoet/FieldSpec.java

@@ -169,6 +169,10 @@ public Builder initializer(CodeBlock codeBlock) {
       return this;
     }
 
+    public Builder initializer(LambdaSpec lambda) {
+      return initializer(lambda.toString());
+    }
+
     public FieldSpec build() {
       return new FieldSpec(this);
     }

src/main/java/com/squareup/javapoet/LambdaSpec.java

@@ -0,0 +1,224 @@
+package com.squareup.javapoet;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import javax.lang.model.element.Modifier;
+import static com.squareup.javapoet.Util.checkArgument;
+import static com.squareup.javapoet.Util.checkNotNull;;
+
+public class LambdaSpec {
+
+    /**
+     * The characters that will be placed *before*
+     * the *arrow* of the lambda.
+    */
+    private final String beforeArrow;
+
+    /** The inputs of the function. */
+    private final List<ParameterSpec> inputs;
+
+    /**
+     * The characters that will be placed *before*
+     * the *body* of the lambda (and before the
+     * curly bracket in such cases).
+    */
+    private final String beforeBody;
+
+    /** The body of the function. */
+    private final CodeBlock body;
+
+    /**
+     * The characters that will be placed *after*
+     * the *body* of the lambda (and after the
+     * curly bracket in such cases).
+    */
+    private final String afterBody;
+
+    /** The visibility status of the input parameters */
+    private boolean visibleTypes;
+
+    private LambdaSpec(Builder builder) {
+        this.inputs = checkNotNull(builder.inputs, "inputs == null");
+        validateInputs(this.inputs);
+        this.body = checkNotNull(builder.body, "body == null");
+        this.beforeArrow = checkNotNull(builder.beforeArrow, "beforeArrow == null");
+        this.beforeBody = checkNotNull(builder.beforeBody, "beforeBody == null");
+        this.afterBody = checkNotNull(builder.afterBody, "afterBody == null");
+        this.visibleTypes = builder.visibleTypes;
+    }
+
+    /**
+     * Emits the left hand side of the lambda function, meaning
+     * the inputs.
+     * @param codeWriter the writer used to append the constructed
+     * inputs.
+     * @throws IOException
+     */
+    void emitInputs(CodeWriter codeWriter) throws IOException {
+        boolean placeParenthesis = inputs.size() > 1 || visibleTypes || inputs.isEmpty();
+
+        codeWriter.emit(placeParenthesis ? "(" : "");
+
+        int paramsPlaced = 0;
+
+        // the inputs of the lambda (left side)
+        for (ParameterSpec inputParameter : inputs) {
+            if (visibleTypes) {
+                codeWriter.emitAnnotations(inputParameter.annotations, true);
+                codeWriter.emitModifiers(inputParameter.modifiers);
+                codeWriter.emit(inputParameter.type.toString() + " ");
+            }
+            codeWriter.emit(inputParameter.name);
+            if (++paramsPlaced < inputs.size())
+                codeWriter.emit(", ");
+        }
+
+        codeWriter.emit(placeParenthesis ? ")" : "");
+    }
+
+    /**
+     * Emits the right hand side of the lambda function, meaning
+     * the body.
+     * @param codeWriter the writer used to append the constructed
+     * body.
+     * @throws IOException
+     */
+    void emitBody(CodeWriter codeWriter) throws IOException {
+        String bodySide = this.body.toString();
+
+        // true if the function has more than 1 statement in its body
+        boolean multiStatementBody = bodySide.contains(";");
+
+        codeWriter.emit(multiStatementBody ? "{" : "");
+        codeWriter.emit(body);
+        codeWriter.emit(multiStatementBody ? "}" : "");
+    }
+
+    @Override public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null) return false;
+        if (getClass() != o.getClass()) return false;
+        return toString().equals(o.toString());
+    }
+
+    @Override public int hashCode() {
+        return toString().hashCode();
+    }
+
+    @Override public String toString() {
+        StringBuilder out = new StringBuilder();
+        try {
+            CodeWriter codeWriter = new CodeWriter(out);
+            emitInputs(codeWriter);
+            codeWriter.emit(this.beforeArrow);
+            codeWriter.emit("->");
+            codeWriter.emit(this.beforeBody);
+            emitBody(codeWriter);
+            codeWriter.emit(this.afterBody);
+            return out.toString();
+        } catch (IOException e) {
+            throw new AssertionError();
+        }
+    }
+
+    /**
+     * Validates a list of parameters as lambda inputs.
+     * @param inputParameters the parameters of the lambda.
+     */
+    private static void validateInputs(List<ParameterSpec> inputParameters) {
+        for (ParameterSpec inputParameter : inputParameters) {
+            checkArgument(
+                !inputParameter.type.equals(TypeName.VOID),
+                "lambda input parameters cannot be of void type"
+            );
+            checkArgument(
+                inputParameter.modifiers.stream().allMatch(p -> p.equals(Modifier.FINAL)),
+                "lambda input parameters can only be final"
+            );
+        }
+    }
+
+    public static Builder builder(List<ParameterSpec> inputs, CodeBlock body) {
+        return new Builder(inputs, body);
+    }
+
+    public static Builder builder(List<ParameterSpec> inputs, String body) {
+        return builder(inputs, CodeBlock.of(body));
+    }
+
+    public static Builder builder(CodeBlock body) {
+        return new Builder(body);
+    }
+
+    public static Builder builder(String body) {
+        return builder(CodeBlock.of(body));
+    }
+
+    public Builder toBuilder() {
+        return toBuilder(inputs, body);
+    }
+
+    Builder toBuilder(List<ParameterSpec> inputs, CodeBlock body) {
+        Builder builder = new Builder(inputs, body);
+        return builder;
+    }
+
+    public static final class Builder {
+        public String beforeArrow = " ";
+        public List<ParameterSpec> inputs = new ArrayList<>();
+        public String beforeBody = " ";
+        public CodeBlock body;
+        public String afterBody = "";
+        public boolean visibleTypes;
+
+        private Builder(List<ParameterSpec> inputs, CodeBlock body) {
+            this.inputs = inputs;
+            this.body = body;
+        }
+
+        private Builder(List<ParameterSpec> inputs, String body) {
+            this(inputs, CodeBlock.of(body));
+        }
+
+        private Builder(CodeBlock body) {
+            this.body = body;
+        }
+
+        private Builder(String body) {
+            this(CodeBlock.of(body));
+        }
+
+        /** Adds an input to the function. */
+        public Builder addInput(ParameterSpec... parameters) {
+            Collections.addAll(this.inputs, parameters);
+            return this;
+        }
+
+        public Builder beforeArrow(String code) {
+            this.beforeArrow = code;
+            return this;
+        }
+
+        public Builder beforeBody(String code) {
+            this.beforeBody = code;
+            return this;
+        }
+
+        public Builder afterBody(String code) {
+            this.afterBody = code;
+            return this;
+        }
+
+        public Builder visibleTypes() {
+            this.visibleTypes = true;
+            return this;
+        }
+
+        public LambdaSpec build() {
+            return new LambdaSpec(this);
+        }
+    }
+}

src/main/java/com/squareup/javapoet/MethodSpec.java

@@ -466,6 +466,11 @@ public Builder defaultValue(CodeBlock codeBlock) {
       return this;
     }
 
+    public Builder addLambda(LambdaSpec lambda) {
+      code.addLambda(lambda);
+      return this;
+    }
+
     /**
      * @param controlFlow the control flow construct and its code, such as "if (foo == 5)".
      * Shouldn't contain braces or newline characters.

src/test/java/com/squareup/javapoet/LambdaSpecTest.java

@@ -0,0 +1,126 @@
+package com.squareup.javapoet;
+
+import javax.lang.model.element.Modifier;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
+
+import org.junit.Test;
+
+
+public class LambdaSpecTest {
+
+    /** Ensure that void inputs result in exception. */
+    @Test public void ensureInputTypeError() {
+        // parameter with void type - should be rejected
+        ParameterSpec p1 = ParameterSpec.builder(TypeName.VOID, "x").build();
+
+        try {
+            // check that void type will cause errors
+            LambdaSpec.builder("2 * x").addInput(p1).build();
+
+            fail("Managed to create a lambda with void type input.");
+        } catch (IllegalArgumentException expected) {
+            assertThat(expected).hasMessageThat().isEqualTo(
+                "lambda input parameters cannot be of void type"
+            );
+        }
+    }
+
+    /** Ensure that not all type inputs are allowed. */
+    @Test public void ensureInputModifierError() {
+        // parameter with modifier different than final - should be rejected
+        ParameterSpec p1 = ParameterSpec.builder(TypeName.INT, "x")
+            .addModifiers(Modifier.PUBLIC).build();
+
+        try {
+            // check that input will cause errors
+            LambdaSpec.builder("2 * x").addInput(p1).build();
+
+            fail("Managed to create a lambda with input of public access.");
+        } catch (IllegalArgumentException expected) {
+            assertThat(expected).hasMessageThat().isEqualTo(
+                "lambda input parameters can only be final"
+            );
+        }
+    }
+
+    @Test public void ensureMethodSpecCompatibility() {
+        // lambda inputs
+        ParameterSpec p1 = ParameterSpec.builder(TypeName.INT, "x").build();
+        ParameterSpec p2 = ParameterSpec.builder(TypeName.DOUBLE, "y").build();
+
+        // lambda body that considers input values
+        CodeBlock body1 = CodeBlock.of("int $3N = 3; return $1N + $2N + $3N;", "x", "y", "z");
+
+        // lambda body that does not consider input values
+        CodeBlock body2 = CodeBlock.of("int $1N = 3; int $2N = 5; return $1N + $2N;", "x", "y");
+
+        // lambdas of different nature
+        LambdaSpec lambda1 = LambdaSpec.builder(body1).addInput(p1, p2).visibleTypes().build();
+        LambdaSpec lambda2 = LambdaSpec.builder("x + y").addInput(p2).build();
+        LambdaSpec lambda3 = LambdaSpec.builder(body2).build();
+        LambdaSpec lambda4 = LambdaSpec.builder("5 + 7").build();
+        LambdaSpec lambda5 = LambdaSpec.builder("method1(); method2();").build();
+        LambdaSpec lambda6 = LambdaSpec.builder("x + 5").addInput(p1).visibleTypes().build();
+
+        MethodSpec method = MethodSpec.methodBuilder("method")
+        .addCode("methodCall(")
+        .addLambda(lambda1).addCode(", ") // lambda with multiple inputs and (CodeBlock) body and visible types
+        .addLambda(lambda2).addCode(", ") // lambda with single input and (String) body
+        .addLambda(lambda3).addCode(", ") // lambda with no inputs and (CodeBlock) body
+        .addLambda(lambda4).addCode(", ") // lambda with no inputs and (String) body
+        .addLambda(lambda5).addCode(", ") // lambda with multiple statements
+        .addLambda(lambda6).addCode(");\n") // lambda with single input of visible type
+        .build();
+
+        assertThat(method.toString()).isEqualTo(""
+          + "void method() {\n"
+          +  "  methodCall("
+          +    "(int x, double y) -> {int z = 3; return x + y + z;}, "
+          +    "y -> x + y, "
+          +    "() -> {int x = 3; int y = 5; return x + y;}, "
+          +    "() -> 5 + 7, "
+          +    "() -> {method1(); method2();}, "
+          +    "(int x) -> x + 5"
+          +  ");\n"
+          + "}\n"
+        );
+    }
+
+    @Test public void ensureCodeBlockCompatibility() {
+        // lambda body that does not consider input values
+        CodeBlock body = CodeBlock.of("int $1N = 3; int $2N = 5; return $1N + $2N;", "x", "y");
+
+        // lambda expression that does not consider input values
+        CodeBlock body2 = CodeBlock.of("5 + 3");
+
+        LambdaSpec lambda1 = LambdaSpec.builder(body).build();
+        LambdaSpec lambda2 = LambdaSpec.builder(body2).build();
+        LambdaSpec lambda3 = LambdaSpec.builder(body).beforeBody("\n  ").build();
+
+        CodeBlock codeWithLambda = CodeBlock.builder()
+        .add("methodCall(")
+        .addLambda(lambda1).add(", ") // producer lambda
+        .addLambda(lambda2).add(");\n")
+        .add("Producer<Integer> pr = ")
+        .addLambda(lambda3)
+        .build();
+
+        MethodSpec method = MethodSpec.methodBuilder("method")
+        .addCode(codeWithLambda)
+        .build();
+
+        assertThat(method.toString()).isEqualTo(""
+        + "void method() {\n"
+        +  "  methodCall("
+        +  "() -> {int x = 3; int y = 5; return x + y;}, "
+        +   "() -> 5 + 3"
+        +  ");\n"
+        +  "  Producer<Integer> pr = () ->\n"
+        +  "    {int x = 3; int y = 5; return x + y;}\n"
+        +  "}\n"
+      );
+    }
+
+}