MethodSpec factory for overriding an ExecutableElement.

square · Feb 16, 2015 · 885ab84 · 885ab84
1 parent 33bb07e
commit 885ab84
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 10 deletions.
diff --git a/src/main/java/com/squareup/javapoet/MethodSpec.java b/src/main/java/com/squareup/javapoet/MethodSpec.java
@@ -19,13 +19,19 @@
 import java.io.StringWriter;
 import java.lang.reflect.Type;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Iterator;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 import javax.lang.model.SourceVersion;
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeParameterElement;
+import javax.lang.model.element.VariableElement;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.type.TypeVariable;
 
 import static com.squareup.javapoet.Util.checkArgument;
 import static com.squareup.javapoet.Util.checkNotNull;
@@ -150,6 +156,42 @@ public static Builder constructorBuilder() {
     return new Builder(CONSTRUCTOR);
   }
 
+  public static Builder overriding(ExecutableElement method) {
+    checkNotNull(method, "method == null");
+
+    Set<Modifier> modifiers = method.getModifiers();
+    if (modifiers.contains(Modifier.PRIVATE)
+        || modifiers.contains(Modifier.FINAL)
+        || modifiers.contains(Modifier.STATIC)) {
+      throw new IllegalArgumentException("cannot override method with modifiers: " + modifiers);
+    }
+
+    String methodName = method.getSimpleName().toString();
+    MethodSpec.Builder methodBuilder = MethodSpec.methodBuilder(methodName);
+    methodBuilder.addAnnotation(Override.class);
+    methodBuilder.returns(TypeName.get(method.getReturnType()));
+
+    modifiers = new LinkedHashSet<>(modifiers); // Local copy so we can remove.
+    modifiers.remove(Modifier.ABSTRACT);
+    methodBuilder.addModifiers(modifiers.toArray(new Modifier[modifiers.size()]));
+
+    for (VariableElement parameter : method.getParameters()) {
+      methodBuilder.addParameter(TypeName.get(parameter.asType()),
+          parameter.getSimpleName().toString());
+    }
+
+    for (TypeMirror thrownType : method.getThrownTypes()) {
+      methodBuilder.addException(TypeName.get(thrownType));
+    }
+
+    for (TypeParameterElement typeParameterElement : method.getTypeParameters()) {
+      methodBuilder.addTypeVariable(
+          TypeVariableName.get((TypeVariable) typeParameterElement.asType()));
+    }
+
+    return methodBuilder;
+  }
+
   public static final class Builder {
     private final String name;
 

diff --git a/src/test/java/com/squareup/javapoet/MethodSpecTest.java b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
@@ -15,10 +15,15 @@
  */
 package com.squareup.javapoet;
 
+import com.google.common.collect.ImmutableSet;
+import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.Modifier;
 import org.junit.Test;
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public final class MethodSpecTest {
 
@@ -27,8 +32,7 @@ public final class MethodSpecTest {
       MethodSpec.methodBuilder("doSomething").addAnnotations(null);
       fail();
     } catch (IllegalArgumentException expected) {
-      assertThat(expected.getMessage())
-          .isEqualTo("annotationSpecs == null");
+      assertThat(expected).hasMessage("annotationSpecs == null");
     }
   }
 
@@ -37,8 +41,7 @@ public final class MethodSpecTest {
       MethodSpec.methodBuilder("doSomething").addTypeVariables(null);
       fail();
     } catch (IllegalArgumentException expected) {
-      assertThat(expected.getMessage())
-          .isEqualTo("typeVariables == null");
+      assertThat(expected).hasMessage("typeVariables == null");
     }
   }
 
@@ -47,8 +50,7 @@ public final class MethodSpecTest {
       MethodSpec.methodBuilder("doSomething").addParameters(null);
       fail();
     } catch (IllegalArgumentException expected) {
-      assertThat(expected.getMessage())
-          .isEqualTo("parameterSpecs == null");
+      assertThat(expected).hasMessage("parameterSpecs == null");
     }
   }
 
@@ -57,8 +59,32 @@ public final class MethodSpecTest {
       MethodSpec.methodBuilder("doSomething").addExceptions(null);
       fail();
     } catch (IllegalArgumentException expected) {
-      assertThat(expected.getMessage())
-          .isEqualTo("exceptions == null");
+      assertThat(expected).hasMessage("exceptions == null");
     }
   }
-}
+
+  @Test public void overrideInvalidModifiers() {
+    ExecutableElement method = mock(ExecutableElement.class);
+    when(method.getModifiers()).thenReturn(ImmutableSet.of(Modifier.FINAL));
+    try {
+      MethodSpec.overriding(method);
+      fail();
+    } catch (IllegalArgumentException expected) {
+      assertThat(expected).hasMessage("cannot override method with modifiers: [final]");
+    }
+    when(method.getModifiers()).thenReturn(ImmutableSet.of(Modifier.PRIVATE));
+    try {
+      MethodSpec.overriding(method);
+      fail();
+    } catch (IllegalArgumentException expected) {
+      assertThat(expected).hasMessage("cannot override method with modifiers: [private]");
+    }
+    when(method.getModifiers()).thenReturn(ImmutableSet.of(Modifier.STATIC));
+    try {
+      MethodSpec.overriding(method);
+      fail();
+    } catch (IllegalArgumentException expected) {
+      assertThat(expected).hasMessage("cannot override method with modifiers: [static]");
+    }
+  }
+}