Merge pull request #644 from square/jwilson.0428.avoid_typeelement_ge…

…tkind Avoid TypeElement.getKind() in ClassName.get().
square · Apr 29, 2018 · 30ff8ce · 30ff8ce
2 parents 3f4d5c2 + 1c897e4
commit 30ff8ce
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 20 deletions.
diff --git a/src/main/java/com/squareup/javapoet/ClassName.java b/src/main/java/com/squareup/javapoet/ClassName.java
@@ -20,15 +20,12 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import javax.lang.model.element.Element;
-import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.Name;
 import javax.lang.model.element.PackageElement;
 import javax.lang.model.element.TypeElement;
 
 import static com.squareup.javapoet.Util.checkArgument;
 import static com.squareup.javapoet.Util.checkNotNull;
-import static javax.lang.model.element.NestingKind.MEMBER;
-import static javax.lang.model.element.NestingKind.TOP_LEVEL;
 
 /** A fully-qualified class name for top-level and member classes. */
 public final class ClassName extends TypeName implements Comparable<ClassName> {
@@ -207,27 +204,20 @@ public static ClassName get(String packageName, String simpleName, String... sim
   /** Returns the class name for {@code element}. */
   public static ClassName get(TypeElement element) {
     checkNotNull(element, "element == null");
-    checkArgument(element.getNestingKind() == TOP_LEVEL || element.getNestingKind() == MEMBER,
-        "unexpected type nesting");
     String simpleName = element.getSimpleName().toString();
 
-    if (isClassOrInterface(element.getEnclosingElement())) {
-      return ClassName.get((TypeElement) element.getEnclosingElement()).nestedClass(simpleName);
-    }
-
-    String packageName = getPackage(element.getEnclosingElement()).getQualifiedName().toString();
-    return new ClassName(packageName, null, simpleName);
-  }
+    switch (element.getNestingKind()) {
+      case TOP_LEVEL:
+        Name packageName = ((PackageElement) element.getEnclosingElement()).getQualifiedName();
+        return new ClassName(packageName.toString(), null, simpleName);
 
-  private static boolean isClassOrInterface(Element e) {
-    return e.getKind().isClass() || e.getKind().isInterface();
-  }
+      case MEMBER:
+        ClassName enclosingClass = get((TypeElement) element.getEnclosingElement());
+        return enclosingClass.nestedClass(simpleName);
 
-  private static PackageElement getPackage(Element type) {
-    while (type.getKind() != ElementKind.PACKAGE) {
-      type = type.getEnclosingElement();
+      default:
+        throw new IllegalArgumentException("unexpected type nesting");
     }
-    return (PackageElement) type;
   }
 
   @Override public int compareTo(ClassName o) {

diff --git a/src/test/java/com/squareup/javapoet/ClassNameTest.java b/src/test/java/com/squareup/javapoet/ClassNameTest.java
@@ -23,10 +23,12 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.when;
 
 @RunWith(JUnit4.class)
 public final class ClassNameTest {
@@ -110,6 +112,36 @@ static class $Inner {}
     assertThat(ClassName.get(inner).toString()).isEqualTo("com.squareup.javapoet.ClassNameTest.$Outer.$Inner");
   }
 
+  /**
+   * Buck builds with "source-based ABI generation" and those builds don't support
+   * {@link TypeElement#getKind()}. Test to confirm that we don't use that API.
+   */
+  @Test public void classNameFromTypeElementDoesntUseGetKind() {
+    Elements elements = compilationRule.getElements();
+    TypeElement object = elements.getTypeElement(Object.class.getCanonicalName());
+    assertThat(ClassName.get(preventGetKind(object)).toString())
+        .isEqualTo("java.lang.Object");
+    TypeElement outer = elements.getTypeElement($Outer.class.getCanonicalName());
+    assertThat(ClassName.get(preventGetKind(outer)).toString())
+        .isEqualTo("com.squareup.javapoet.ClassNameTest.$Outer");
+    TypeElement inner = elements.getTypeElement($Outer.$Inner.class.getCanonicalName());
+    assertThat(ClassName.get(preventGetKind(inner)).toString())
+        .isEqualTo("com.squareup.javapoet.ClassNameTest.$Outer.$Inner");
+  }
+
+  /** Returns a new instance like {@code object} that throws on {@code getKind()}. */
+  private TypeElement preventGetKind(TypeElement object) {
+    TypeElement spy = Mockito.spy(object);
+    when(spy.getKind()).thenThrow(new AssertionError());
+    when(spy.getEnclosingElement()).thenAnswer(invocation -> {
+      Object enclosingElement = invocation.callRealMethod();
+      return enclosingElement instanceof TypeElement
+          ? preventGetKind((TypeElement) enclosingElement)
+          : enclosingElement;
+    });
+    return spy;
+  }
+
   @Test public void classNameFromClass() {
     assertThat(ClassName.get(Object.class).toString())
         .isEqualTo("java.lang.Object");