Avoid TypeElement.getKind() in ClassName.get().

There's some bug that's triggered by Buck's ABI builds and avoiding this
API avoids triggering that bug.

2 files changed, 42 insertions, 20 deletions

diff --git a/src/main/java/com/squareup/javapoet/ClassName.java b/src/main/java/com/squareup/javapoet/ClassName.java
index e8af154..766d9e2 100644
--- 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.ArrayList;
 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 final class ClassName extends TypeName implements Comparable<ClassName> {
   /** 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
index 0488213..e2cc55e 100644
--- 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.Rule;
 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 @@ public final class ClassNameTest {
     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");