Add generic import support for Codegen

Supports any kind of DataClass or NonNull-type annotations instead of
enforcing the Android framework package.

Bug: 195012034

Test: manual, generate for DomainVerificationInfo, verify no change
Test: manual, generate for mainline class, verify no internal import

Change-Id: I0360281aa20b7f2c6e75dec88a92b810fa69b9b7

1 files changed, 18 insertions, 3 deletions

diff --git a/tools/codegen/src/com/android/codegen/ImportsProvider.kt b/tools/codegen/src/com/android/codegen/ImportsProvider.kt
index 27dd9587db25..46df2739e59f 100644
--- a/tools/codegen/src/com/android/codegen/ImportsProvider.kt
+++ b/tools/codegen/src/com/android/codegen/ImportsProvider.kt
@@ -53,10 +53,25 @@ interface ImportsProvider {
      * Optionally shortens a class reference if there's a corresponding import present
      */
     fun classRef(fullName: String): String {
-
         val pkg = fullName.substringBeforeLast(".")
         val simpleName = fullName.substringAfterLast(".")
-        if (fileAst.imports.any { imprt ->
+        val imports = fileAst.imports
+
+        // If an import of the same class name is available,
+        // use it instead of the internal Android package variants.
+        if (fullName.startsWith("com.android.internal.util.")
+            && imports.any {
+                it.nameAsString.endsWith(fullName.removePrefix("com.android.internal.util."))
+            }
+        ) {
+            return fullName.removePrefix("com.android.internal.util.")
+        } else if (fullName.startsWith("android.annotation")
+            && imports.any { it.nameAsString.endsWith(simpleName) }
+        ) {
+            return simpleName
+        }
+
+        if (imports.any { imprt ->
                     imprt.nameAsString == fullName
                             || (imprt.isAsterisk && imprt.nameAsString == pkg)
                 }) {
@@ -89,4 +104,4 @@ interface ImportsProvider {
 /** @see classRef */
 inline fun <reified T : Any> ImportsProvider.classRef(): String {
     return classRef(T::class.java.name)
-}
\ No newline at end of file
+}