Add checks to ParameterSpec with VariableElement + copy over annotations (#681)

* Add checks to ParameterSpec with VariableElement + copy over annotations

Signed-off-by: shaishavgandhi05 <shaishgandhi@gmail.com>

* Add test for variable element

Signed-off-by: shaishavgandhi05 <shaishgandhi@gmail.com>

* Extract util methods into TestUtil

* Fix formatting

* Make findFirst more generic

Co-authored-by: Egor Andreevich <andreevich.egor@gmail.com>

4 files changed, 89 insertions, 12 deletions

diff --git a/src/main/java/com/squareup/javapoet/ParameterSpec.java b/src/main/java/com/squareup/javapoet/ParameterSpec.java
index e98b99e..386ed1a 100644
--- a/src/main/java/com/squareup/javapoet/ParameterSpec.java
+++ b/src/main/java/com/squareup/javapoet/ParameterSpec.java
@@ -21,7 +21,9 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 import javax.lang.model.SourceVersion;
+import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.VariableElement;
@@ -83,10 +85,19 @@ public final class ParameterSpec {
   }
 
   public static ParameterSpec get(VariableElement element) {
+    checkArgument(element.getKind().equals(ElementKind.PARAMETER), "element is not a parameter");
+
+    // Copy over any annotations from element.
+    List<AnnotationSpec> annotations = element.getAnnotationMirrors()
+        .stream()
+        .map((mirror) -> AnnotationSpec.get(mirror))
+        .collect(Collectors.toList());
+
     TypeName type = TypeName.get(element.asType());
     String name = element.getSimpleName().toString();
     return ParameterSpec.builder(type, name)
         .addModifiers(element.getModifiers())
+        .addAnnotations(annotations)
         .build();
   }
 
diff --git a/src/test/java/com/squareup/javapoet/MethodSpecTest.java b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
index b2c7106..3db519c 100644
--- a/src/test/java/com/squareup/javapoet/MethodSpecTest.java
+++ b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
@@ -38,6 +38,7 @@ import org.junit.Test;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static com.squareup.javapoet.MethodSpec.CONSTRUCTOR;
+import static com.squareup.javapoet.TestUtil.findFirst;
 import static javax.lang.model.util.ElementFilter.methodsIn;
 import static org.junit.Assert.fail;
 
@@ -56,15 +57,6 @@ public final class MethodSpecTest {
     return elements.getTypeElement(clazz.getCanonicalName());
   }
 
-  private ExecutableElement findFirst(Collection<ExecutableElement> elements, String name) {
-    for (ExecutableElement executableElement : elements) {
-      if (executableElement.getSimpleName().toString().equals(name)) {
-        return executableElement;
-      }
-    }
-    throw new IllegalArgumentException(name + " not found in " + elements);
-  }
-
   @Test public void nullAnnotationsAddition() {
     try {
       MethodSpec.methodBuilder("doSomething").addAnnotations(null);
@@ -155,8 +147,8 @@ public final class MethodSpecTest {
         + "@java.lang.Override\n"
         + "protected <T extends java.lang.Runnable & java.io.Closeable> java.lang.Runnable "
         + "everything(\n"
-        + "    java.lang.String arg0, java.util.List<? extends T> arg1) throws java.io.IOException,\n"
-        + "    java.lang.SecurityException {\n"
+        + "    @com.squareup.javapoet.MethodSpecTest.Nullable java.lang.String arg0,\n"
+        + "    java.util.List<? extends T> arg1) throws java.io.IOException, java.lang.SecurityException {\n"
         + "}\n");
   }
 
@@ -187,7 +179,9 @@ public final class MethodSpecTest {
     TypeElement classElement = getElement(ExtendsIterableWithDefaultMethods.class);
     DeclaredType classType = (DeclaredType) classElement.asType();
     List<ExecutableElement> methods = methodsIn(elements.getAllMembers(classElement));
-    ExecutableElement exec = findFirst(methods, "spliterator");
+    ExecutableElement exec = 
+      
+      (methods, "spliterator");
     MethodSpec method = MethodSpec.overriding(exec, classType, types).build();
     assertThat(method.toString()).isEqualTo(""
         + "@java.lang.Override\n"
diff --git a/src/test/java/com/squareup/javapoet/ParameterSpecTest.java b/src/test/java/com/squareup/javapoet/ParameterSpecTest.java
index f66850b..954db39 100644
--- a/src/test/java/com/squareup/javapoet/ParameterSpecTest.java
+++ b/src/test/java/com/squareup/javapoet/ParameterSpecTest.java
@@ -15,16 +15,39 @@
  */
 package com.squareup.javapoet;
 
+import com.google.testing.compile.CompilationRule;
 import java.util.ArrayList;
 import java.util.List;
+import javax.annotation.Nullable;
+import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.VariableElement;
+import javax.lang.model.util.Elements;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.squareup.javapoet.TestUtil.findFirst;
+import static javax.lang.model.util.ElementFilter.fieldsIn;
+import static javax.lang.model.util.ElementFilter.methodsIn;
 import static org.junit.Assert.fail;
 
 import javax.lang.model.element.Modifier;
 
 public class ParameterSpecTest {
+  @Rule public final CompilationRule compilation = new CompilationRule();
+
+  private Elements elements;
+
+  @Before public void setUp() {
+    elements = compilation.getElements();
+  }
+
+  private TypeElement getElement(Class<?> clazz) {
+    return elements.getTypeElement(clazz.getCanonicalName());
+  }
+
   @Test public void equalsAndHashCode() {
     ParameterSpec a = ParameterSpec.builder(int.class, "foo").build();
     ParameterSpec b = ParameterSpec.builder(int.class, "foo").build();
@@ -48,6 +71,38 @@ public class ParameterSpecTest {
     }
   }
 
+  final class VariableElementFieldClass {
+    String name;
+  }
+
+  @Test public void fieldVariableElement() {
+    TypeElement classElement = getElement(VariableElementFieldClass.class);
+    List<VariableElement> methods = fieldsIn(elements.getAllMembers(classElement));
+    VariableElement element = findFirst(methods, "name");
+
+    try {
+      ParameterSpec.get(element);
+      fail();
+    } catch (IllegalArgumentException exception) {
+      assertThat(exception).hasMessageThat().isEqualTo("element is not a parameter");
+    }
+  }
+
+  final class VariableElementParameterClass {
+    public void foo(@Nullable final String bar) {
+    }
+  }
+
+  @Test public void parameterVariableElement() {
+    TypeElement classElement = getElement(VariableElementParameterClass.class);
+    List<ExecutableElement> methods = methodsIn(elements.getAllMembers(classElement));
+    ExecutableElement element = findFirst(methods, "foo");
+    VariableElement parameterElement = element.getParameters().get(0);
+
+    assertThat(ParameterSpec.get(parameterElement).toString())
+        .isEqualTo("@javax.annotation.Nullable java.lang.String arg0");
+  }
+
   @Test public void addNonFinalModifier() {
     List<Modifier> modifiers = new ArrayList<>();
     modifiers.add(Modifier.FINAL);
diff --git a/src/test/java/com/squareup/javapoet/TestUtil.java b/src/test/java/com/squareup/javapoet/TestUtil.java
new file mode 100644
index 0000000..f773d50
--- /dev/null
+++ b/src/test/java/com/squareup/javapoet/TestUtil.java
@@ -0,0 +1,17 @@
+package com.squareup.javapoet;
+
+import javax.lang.model.element.Element;
+import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.VariableElement;
+import java.util.Collection;
+
+final class TestUtil {
+  static <E extends Element> E findFirst(Collection<E> elements, String name) {
+    for (E element : elements) {
+      if (element.getSimpleName().toString().equals(name)) {
+        return element;
+      }
+    }
+    throw new IllegalArgumentException(name + " not found in " + elements);
+  }
+}