Extensions to convert JavaPoet types to KotlinPoet and vice-versa

JavaPoetExt.kt

import com.squareup.javapoet.ClassName
import com.squareup.javapoet.ParameterizedTypeName
import com.squareup.javapoet.TypeName
import com.squareup.javapoet.TypeVariableName
import com.squareup.javapoet.WildcardTypeName
import com.squareup.javapoet.ArrayTypeName
import com.squareup.kotlinpoet.ARRAY
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy

fun ClassName.asKPClassName(): com.squareup.kotlinpoet.ClassName {
    return if (simpleNames().size == 1) {
        com.squareup.kotlinpoet.ClassName(packageName(), simpleName())
    } else {
        com.squareup.kotlinpoet.ClassName(packageName(), simpleNames().first(), *simpleNames().drop(1).toTypedArray())
    }
}

fun ParameterizedTypeName.asKPParameterizedTypeName(): com.squareup.kotlinpoet.ParameterizedTypeName {
    return rawType.asKPClassName().parameterizedBy(*typeArguments.map { it.asKP() }.toTypedArray())
}

fun TypeVariableName.asKPTypeVariableName(): com.squareup.kotlinpoet.TypeVariableName {
    return if (bounds.isEmpty()) {
        com.squareup.kotlinpoet.TypeVariableName(name)
    } else {
        com.squareup.kotlinpoet.TypeVariableName(name, *bounds.map { it.asKP() }.toTypedArray())
    }
}

fun TypeName.asKP(): com.squareup.kotlinpoet.TypeName {
    return when (this) {
        is ClassName -> when (this) {
            TypeName.BOOLEAN.box() -> com.squareup.kotlinpoet.BOOLEAN
            TypeName.BYTE.box() -> com.squareup.kotlinpoet.BYTE
            TypeName.CHAR.box() -> com.squareup.kotlinpoet.CHAR
            TypeName.SHORT.box() -> com.squareup.kotlinpoet.SHORT
            TypeName.INT.box() -> com.squareup.kotlinpoet.INT
            TypeName.LONG.box() -> com.squareup.kotlinpoet.LONG
            TypeName.FLOAT.box() -> com.squareup.kotlinpoet.FLOAT
            TypeName.DOUBLE.box() -> com.squareup.kotlinpoet.DOUBLE
            TypeName.OBJECT -> com.squareup.kotlinpoet.ANY
            PoetInterop.CN_JAVA_STRING -> PoetInterop.CN_KOTLIN_STRING
            PoetInterop.CN_JAVA_LIST -> PoetInterop.CN_KOTLIN_LIST
            PoetInterop.CN_JAVA_SET -> PoetInterop.CN_KOTLIN_SET
            PoetInterop.CN_JAVA_MAP -> PoetInterop.CN_KOTLIN_MAP
            else -> asKPClassName()
        }
        is ParameterizedTypeName -> asKPParameterizedTypeName()
        is TypeVariableName -> asKPTypeVariableName()
        is WildcardTypeName -> kotlin.TODO()
        is ArrayTypeName -> ARRAY.parameterizedBy(componentType.asKP())
        else -> when (unboxIfBoxedPrimitive()) {
            TypeName.BOOLEAN -> com.squareup.kotlinpoet.BOOLEAN
            TypeName.BYTE -> com.squareup.kotlinpoet.BYTE
            TypeName.CHAR -> com.squareup.kotlinpoet.CHAR
            TypeName.SHORT -> com.squareup.kotlinpoet.SHORT
            TypeName.INT -> com.squareup.kotlinpoet.INT
            TypeName.LONG -> com.squareup.kotlinpoet.LONG
            TypeName.FLOAT -> com.squareup.kotlinpoet.FLOAT
            TypeName.DOUBLE -> com.squareup.kotlinpoet.DOUBLE
            else -> kotlin.TODO("Unrecognized type $this")
        }
    }
}

fun TypeName.unboxIfBoxedPrimitive(): TypeName {
    return if (isBoxedPrimitive) {
        unbox()
    } else this
}

fun TypeName.boxIfPrimitive(extraCondition: Boolean = true): TypeName {
    return if (extraCondition && isPrimitive && !isBoxedPrimitive) {
        box()
    } else this
}

KotlinPoetExt.kt

import com.squareup.javapoet.ArrayTypeName
import com.squareup.kotlinpoet.ANY
import com.squareup.kotlinpoet.ARRAY
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.BYTE
import com.squareup.kotlinpoet.CHAR
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.DOUBLE
import com.squareup.kotlinpoet.Dynamic
import com.squareup.kotlinpoet.FLOAT
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.LambdaTypeName
import com.squareup.kotlinpoet.ParameterizedTypeName
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.TypeName
import com.squareup.kotlinpoet.TypeVariableName
import com.squareup.kotlinpoet.WildcardTypeName

fun ClassName.asJPClassName(shouldBox: Boolean = false): com.squareup.javapoet.TypeName {
    return when (copy(nullable = false)) {
        BOOLEAN -> com.squareup.javapoet.TypeName.BOOLEAN.boxIfPrimitive(shouldBox || isNullable)
        BYTE -> com.squareup.javapoet.TypeName.BYTE.boxIfPrimitive(shouldBox || isNullable)
        CHAR -> com.squareup.javapoet.TypeName.CHAR.boxIfPrimitive(shouldBox || isNullable)
        SHORT -> com.squareup.javapoet.TypeName.SHORT.boxIfPrimitive(shouldBox || isNullable)
        INT -> com.squareup.javapoet.TypeName.INT.boxIfPrimitive(shouldBox || isNullable)
        LONG -> com.squareup.javapoet.TypeName.LONG.boxIfPrimitive(shouldBox || isNullable)
        FLOAT -> com.squareup.javapoet.TypeName.FLOAT.boxIfPrimitive(shouldBox || isNullable)
        DOUBLE -> com.squareup.javapoet.TypeName.DOUBLE.boxIfPrimitive(shouldBox || isNullable)
        ANY -> com.squareup.javapoet.TypeName.OBJECT
        PoetInterop.CN_KOTLIN_STRING -> PoetInterop.CN_JAVA_STRING
        PoetInterop.CN_KOTLIN_LIST -> PoetInterop.CN_JAVA_LIST
        PoetInterop.CN_KOTLIN_SET -> PoetInterop.CN_JAVA_SET
        PoetInterop.CN_KOTLIN_MAP -> PoetInterop.CN_JAVA_MAP
        else -> {
            if (simpleNames.size == 1) {
                com.squareup.javapoet.ClassName.get(packageName, simpleName)
            } else {
                com.squareup.javapoet.ClassName.get(packageName, simpleNames.first(), *simpleNames.drop(1).toTypedArray())
            }
        }
    }
}

fun ParameterizedTypeName.asJPParameterizedOrArrayTypeName(): com.squareup.javapoet.TypeName {
    return when (rawType) {
        ARRAY -> {
            val componentType = typeArguments.firstOrNull()?.asJP()
                ?: throw IllegalStateException("Array with no type! $this")
            ArrayTypeName.of(componentType)
        }
        else -> {
            com.squareup.javapoet.ParameterizedTypeName.get(rawType.asJPClassName() as com.squareup.javapoet.ClassName,
                *typeArguments.map { it.asJP(shouldBox = true) }.toTypedArray())
        }
    }
}

fun ParameterizedTypeName.asJPParameterizedTypeName(): com.squareup.javapoet.ParameterizedTypeName {
    check(rawType != ARRAY) {
        "Array type! JavaPoet arrays are a custom TypeName. Use this function only for things you know are not arrays"
    }
    return asJPParameterizedOrArrayTypeName() as com.squareup.javapoet.ParameterizedTypeName
}

fun TypeVariableName.asJPTypeVariableName(): com.squareup.javapoet.TypeVariableName {
    return com.squareup.javapoet.TypeVariableName.get(name, *bounds.map { it.asJP(shouldBox = true) }.toTypedArray())
}

fun TypeName.asJP(shouldBox: Boolean = false): com.squareup.javapoet.TypeName {
    return when (this) {
        is ClassName -> asJPClassName(shouldBox)
        Dynamic -> throw IllegalStateException("Not applicable in Java!")
        is LambdaTypeName -> throw IllegalStateException("Not applicable in Java!")
        is ParameterizedTypeName -> asJPParameterizedOrArrayTypeName()
        is TypeVariableName -> asJPTypeVariableName()
        is WildcardTypeName -> TODO()
    }
}

PoetInterop.kt

/** Various JavaPoet and KotlinPoet representations of some common types. */
internal object PoetInterop {
    internal val CN_KOTLIN_STRING = com.squareup.kotlinpoet.ClassName("kotlin", "String")
    internal val CN_KOTLIN_LIST = com.squareup.kotlinpoet.ClassName("kotlin", "List")
    internal val CN_KOTLIN_SET = com.squareup.kotlinpoet.ClassName("kotlin", "Set")
    internal val CN_KOTLIN_MAP = com.squareup.kotlinpoet.ClassName("kotlin", "Map")
    internal val CN_JAVA_STRING = com.squareup.javapoet.ClassName.get("java.lang", "String")
    internal val CN_JAVA_LIST = com.squareup.javapoet.ClassName.get("java.util", "List")
    internal val CN_JAVA_SET = com.squareup.javapoet.ClassName.get("java.util", "Set")
    internal val CN_JAVA_MAP = com.squareup.javapoet.ClassName.get("java.util", "Map")
}

PoetInteropTest.kt

import com.google.common.truth.Truth.assertThat
import com.squareup.javapoet.ArrayTypeName
import com.squareup.javapoet.ClassName
import com.squareup.javapoet.TypeName
import com.squareup.kotlinpoet.ANY
import com.squareup.kotlinpoet.ARRAY
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.BYTE
import com.squareup.kotlinpoet.CHAR
import com.squareup.kotlinpoet.DOUBLE
import com.squareup.kotlinpoet.FLOAT
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.asClassName
import com.ubercab.test.UberTestBase
import org.junit.Test

class PoetInteropTest : UberTestBase() {

  @Test
  fun interopClassNamesMatch() {
    val kotlinPoetCN = PoetInteropTest::class.asClassName()
    val javapoetCN = kotlinPoetCN.asJPClassName()

    assertThat(javapoetCN.asKP()).isEqualTo(kotlinPoetCN)
    assertThat(ClassName.get(PoetInteropTest::class.java)).isEqualTo(javapoetCN)
  }

  @Test
  fun interopNestedClassNamesMatch() {
    val kotlinPoetCN = PoetInteropTest::class.asClassName().nestedClass("Foo").nestedClass("Bar")
    val javapoetCN = kotlinPoetCN.asJPClassName()

    assertThat(javapoetCN.asKP()).isEqualTo(kotlinPoetCN)
    assertThat(ClassName.get(PoetInteropTest::class.java).nestedClass("Foo").nestedClass("Bar"))
        .isEqualTo(javapoetCN)
  }

  @Test
  fun kotlinBuiltinsMapCorrectlyToJava() {
    assertThat(PoetInterop.CN_KOTLIN_LIST.asJP()).isEqualTo(PoetInterop.CN_JAVA_LIST)
    assertThat(PoetInterop.CN_KOTLIN_SET.asJP()).isEqualTo(PoetInterop.CN_JAVA_SET)
    assertThat(PoetInterop.CN_KOTLIN_MAP.asJP()).isEqualTo(PoetInterop.CN_JAVA_MAP)
    assertThat(PoetInterop.CN_KOTLIN_STRING.asJP()).isEqualTo(PoetInterop.CN_JAVA_STRING)
    assertThat(ANY.asJP()).isEqualTo(TypeName.OBJECT)
    assertThat(PoetInterop.CN_JAVA_LIST.asKP()).isEqualTo(PoetInterop.CN_KOTLIN_LIST)
    assertThat(PoetInterop.CN_JAVA_SET.asKP()).isEqualTo(PoetInterop.CN_KOTLIN_SET)
    assertThat(PoetInterop.CN_JAVA_MAP.asKP()).isEqualTo(PoetInterop.CN_KOTLIN_MAP)
    assertThat(PoetInterop.CN_JAVA_STRING.asKP()).isEqualTo(PoetInterop.CN_KOTLIN_STRING)
    assertThat(TypeName.OBJECT.asKP()).isEqualTo(ANY)
    // There are more we may add in the future, but these are all needed for modelgen right now
    // See https://github.com/square/kotlinpoet/pull/685
  }

  @Test
  fun boxRequestYieldsBoxedPrimitive() {
    assertThat(BOOLEAN.asJP(shouldBox = true)).isEqualTo(TypeName.BOOLEAN.box())
    assertThat(BYTE.asJP(shouldBox = true)).isEqualTo(TypeName.BYTE.box())
    assertThat(CHAR.asJP(shouldBox = true)).isEqualTo(TypeName.CHAR.box())
    assertThat(SHORT.asJP(shouldBox = true)).isEqualTo(TypeName.SHORT.box())
    assertThat(INT.asJP(shouldBox = true)).isEqualTo(TypeName.INT.box())
    assertThat(LONG.asJP(shouldBox = true)).isEqualTo(TypeName.LONG.box())
    assertThat(FLOAT.asJP(shouldBox = true)).isEqualTo(TypeName.FLOAT.box())
    assertThat(DOUBLE.asJP(shouldBox = true)).isEqualTo(TypeName.DOUBLE.box())
  }

  @Test
  fun primitiveUnboxedByDefault() {
    assertThat(BOOLEAN.asJP()).isEqualTo(TypeName.BOOLEAN)
    assertThat(BYTE.asJP()).isEqualTo(TypeName.BYTE)
    assertThat(CHAR.asJP()).isEqualTo(TypeName.CHAR)
    assertThat(SHORT.asJP()).isEqualTo(TypeName.SHORT)
    assertThat(INT.asJP()).isEqualTo(TypeName.INT)
    assertThat(LONG.asJP()).isEqualTo(TypeName.LONG)
    assertThat(FLOAT.asJP()).isEqualTo(TypeName.FLOAT)
    assertThat(DOUBLE.asJP()).isEqualTo(TypeName.DOUBLE)
  }

  @Test
  fun nullablePrimitiveBoxedByDefault() {
    assertThat(BOOLEAN.copy(nullable = true).asJP()).isEqualTo(TypeName.BOOLEAN.box())
    assertThat(BYTE.copy(nullable = true).asJP()).isEqualTo(TypeName.BYTE.box())
    assertThat(CHAR.copy(nullable = true).asJP()).isEqualTo(TypeName.CHAR.box())
    assertThat(SHORT.copy(nullable = true).asJP()).isEqualTo(TypeName.SHORT.box())
    assertThat(INT.copy(nullable = true).asJP()).isEqualTo(TypeName.INT.box())
    assertThat(LONG.copy(nullable = true).asJP()).isEqualTo(TypeName.LONG.box())
    assertThat(FLOAT.copy(nullable = true).asJP()).isEqualTo(TypeName.FLOAT.box())
    assertThat(DOUBLE.copy(nullable = true).asJP()).isEqualTo(TypeName.DOUBLE.box())
  }

  @Test
  fun arrayTypesConversion() {
    assertThat(ARRAY.parameterizedBy(INT).asJPParameterizedOrArrayTypeName())
        .isEqualTo(ArrayTypeName.of(TypeName.INT))
    assertThat(ARRAY.parameterizedBy(INT.copy(nullable = true)).asJPParameterizedOrArrayTypeName())
        .isEqualTo(ArrayTypeName.of(TypeName.INT.box()))
    assertThat(ArrayTypeName.of(TypeName.INT).asKP()).isEqualTo(ARRAY.parameterizedBy(INT))
    assertThat(ArrayTypeName.of(TypeName.INT.box()).asKP()).isEqualTo(ARRAY.parameterizedBy(INT))
  }
}