Add a tap method to builders

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

@@ -156,7 +156,7 @@ public Builder toBuilder() {
     return builder;
   }
 
-  public static final class Builder {
+  public static final class Builder implements Tapable<Builder> {
     final List<String> formatParts = new ArrayList<>();
     final List<Object> args = new ArrayList<>();
 

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

@@ -106,7 +106,7 @@ public Builder toBuilder() {
     return builder;
   }
 
-  public static final class Builder {
+  public static final class Builder implements Tapable<Builder> {
     private final TypeName type;
     private final String name;
 

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

@@ -291,7 +291,7 @@ public Builder toBuilder() {
     return builder;
   }
 
-  public static final class Builder {
+  public static final class Builder implements Tapable<Builder> {
     private String name;
 
     private final CodeBlock.Builder javadoc = CodeBlock.builder();

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

@@ -134,7 +134,7 @@ Builder toBuilder(TypeName type, String name) {
     return builder;
   }
 
-  public static final class Builder {
+  public static final class Builder implements Tapable<Builder> {
     private final TypeName type;
     private final String name;
     private final CodeBlock.Builder javadoc = CodeBlock.builder();

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2003-2012 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.squareup.javapoet;
+
+import java.util.function.Consumer;
+
+public interface Tapable<T extends Tapable<T>> {
+    @SuppressWarnings("unchecked")
+    default T tap(Consumer<? super T> consumer) {
+        T self = (T) this;
+        consumer.accept(self);
+        return self;
+    }
+}

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

@@ -408,7 +408,7 @@ public enum Kind {
     }
   }
 
-  public static final class Builder {
+  public static final class Builder implements Tapable<Builder> {
     private final Kind kind;
     private final String name;
     private final CodeBlock anonymousTypeArguments;

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

@@ -348,4 +348,19 @@ public final class CodeBlockTest {
 
     assertThat(block.toString()).isEmpty();
   }
+
+  @Test public void canTapInto() {
+    CodeBlock block = CodeBlock.builder()
+            .addStatement("int x = 0")
+            .tap(b -> {
+              int value = 666;
+              b.addStatement("int y = $L", value);
+            })
+            .build();
+    assertThat(block.toString()).isEqualTo(""
+      + "int x = 0;\n"
+      + "int y = 666;\n"
+    );
+
+  }
 }

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

@@ -63,4 +63,12 @@ public class FieldSpecTest {
     builder.modifiers.remove(1);
     assertThat(builder.build().modifiers).containsExactly(Modifier.PUBLIC);
   }
+
+  @Test public void canTapInto() {
+    FieldSpec.Builder builder = FieldSpec.builder(int.class, "foo")
+            .tap(f -> f.addModifiers(Modifier.PUBLIC, Modifier.STATIC));
+
+    builder.modifiers.remove(1);
+    assertThat(builder.build().modifiers).containsExactly(Modifier.PUBLIC);
+  }
 }

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

@@ -481,6 +481,38 @@ abstract static class AbstractClassWithPrivateAnnotation {
         "}\n");
   }
 
+  @Test public void canTapIntoMethod() {
+    MethodSpec spec = MethodSpec.methodBuilder("method")
+            .tap(body -> {
+              for (int i = 0; i<2; i++) {
+                body.addStatement("System.out.println($S)", "Hello " + i);
+              }
+            })
+            .build();
+
+    assertThat(spec.toString()).isEqualTo(""
+            + "void method() {\n"
+            + "  System.out.println(\"Hello 0\");\n"
+            + "  System.out.println(\"Hello 1\");\n"
+            + "}\n");
+  }
+
+  @Test public void canTapIntoConstructor() {
+    MethodSpec spec = MethodSpec.constructorBuilder()
+            .tap(body -> {
+              for (int i = 0; i<2; i++) {
+                body.addStatement("System.out.println($S)", "Hello " + i);
+              }
+            })
+            .build();
+
+    assertThat(spec.toString()).isEqualTo(""
+            + "Constructor() {\n"
+            + "  System.out.println(\"Hello 0\");\n"
+            + "  System.out.println(\"Hello 1\");\n"
+            + "}\n");
+  }
+
   private static CodeBlock named(String format, Map<String, ?> args){
     return CodeBlock.builder().addNamed(format, args).build();
   }

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

@@ -151,4 +151,12 @@ public void foo(@Nullable final String bar) {
     builder.modifiers.remove(1);
     assertThat(builder.build().modifiers).containsExactly(Modifier.PUBLIC);
   }
+
+  @Test public void canTapInto() {
+    ParameterSpec.Builder builder = ParameterSpec.builder(int.class, "foo")
+            .tap(p -> p.addModifiers(Modifier.PUBLIC, Modifier.STATIC));
+
+    builder.modifiers.remove(1);
+    assertThat(builder.build().modifiers).containsExactly(Modifier.PUBLIC);
+  }
 }

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

@@ -2604,7 +2604,7 @@ public void modifyOriginatingElements() {
     builder.originatingElements.clear();
     assertThat(builder.build().originatingElements).isEmpty();
   }
-    
+
   @Test public void javadocWithTrailingLineDoesNotAddAnother() {
     TypeSpec spec = TypeSpec.classBuilder("Taco")
         .addJavadoc("Some doc with a newline\n")
@@ -2634,4 +2634,24 @@ public void modifyOriginatingElements() {
         + "class Taco {\n"
         + "}\n");
   }
+
+  @Test public void canTapInto() {
+    TypeSpec spec = TypeSpec.classBuilder("Taco")
+            .tap(b -> {
+              for (int i = 0; i<2; i++) {
+                b.addMethod(MethodSpec.methodBuilder("foo" + i).build());
+              }
+            })
+            .build();
+
+      assertThat(toString(spec)).isEqualTo(""
+              + "package com.squareup.tacos;\n"
+              + "\n"
+              + "class Taco {\n"
+              + "  void foo0() {\n"
+              + "  }\n\n"
+              + "  void foo1() {\n"
+              + "  }\n"
+              + "}\n");
+  }
 }