Fix extension function imports (#1814)

* Generate imports even for names clashing with locally declared names, but with alias * Force import generation for extension functions - there is no way how an extension function could be used without an import * Add unit tests for the fixed scenarios
square · Jan 29, 2024 · b28e315 · b28e315
1 parent 36d798d
commit b28e315
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 3 deletions.
diff --git a/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/CodeWriter.kt b/kotlinpoet/src/commonMain/kotlin/com/squareup/kotlinpoet/CodeWriter.kt
@@ -517,7 +517,7 @@ internal class CodeWriter constructor(
     if (memberName.packageName.isNotEmpty()) {
       val simpleName = imports[memberName.canonicalName]?.alias ?: memberName.simpleName
       // Check for name clashes with types.
-      if (simpleName !in importableTypes) {
+      if (memberName.isExtension || simpleName !in importableTypes) {
         importableMembers[simpleName] = importableMembers.getValue(simpleName) + memberName
       }
     }
@@ -667,7 +667,7 @@ internal class CodeWriter constructor(
    * collisions, import aliases will be generated.
    */
   private fun suggestedMemberImports(): Map<String, Set<MemberName>> {
-    return importableMembers.filterKeys { it !in referencedNames }.mapValues { it.value.toSet() }
+    return importableMembers.mapValues { it.value.toSet() }
   }
 
   /**
@@ -713,12 +713,14 @@ internal class CodeWriter constructor(
           generatedImports,
           canonicalName = ClassName::canonicalName,
           capitalizeAliases = true,
+          referencedNames = importsCollector.referencedNames,
         )
       val suggestedMemberImports = importsCollector.suggestedMemberImports()
         .generateImports(
           generatedImports,
           canonicalName = MemberName::canonicalName,
           capitalizeAliases = false,
+          referencedNames = importsCollector.referencedNames,
         )
       importsCollector.close()
 
@@ -735,9 +737,10 @@ internal class CodeWriter constructor(
       generatedImports: MutableMap<String, Import>,
       canonicalName: T.() -> String,
       capitalizeAliases: Boolean,
+      referencedNames: Set<String>,
     ): Map<String, T> {
       return flatMap { (simpleName, qualifiedNames) ->
-        if (qualifiedNames.size == 1) {
+        if (qualifiedNames.size == 1 && simpleName !in referencedNames) {
           listOf(simpleName to qualifiedNames.first()).also {
             val canonicalName = qualifiedNames.first().canonicalName()
             generatedImports[canonicalName] = Import(canonicalName)

diff --git a/kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/KotlinPoetTest.kt b/kotlinpoet/src/commonTest/kotlin/com/squareup/kotlinpoet/KotlinPoetTest.kt
@@ -1383,6 +1383,60 @@ class KotlinPoetTest {
     )
   }
 
+  @Test fun extensionFunctionIsImportedEvenIfTheSameIsUsedAlsoFromTheCurrentPackage() {
+    val kotlinIsNullOrEmpty = MemberName(packageName = "kotlin.text", simpleName = "isNullOrEmpty", isExtension = true)
+    val samePackageIsNullOrEmpty = MemberName(packageName = "com.example", simpleName = "isNullOrEmpty", isExtension = true)
+    val file = FileSpec.builder("com.example", "Test")
+      .addFunction(
+        FunSpec.builder("main")
+          .addStatement("val isFirstNull = null.%M()", kotlinIsNullOrEmpty)
+          .addStatement("val isSecondNull = null.%M()", samePackageIsNullOrEmpty)
+          .build(),
+      )
+      .build()
+    assertThat(file.toString()).isEqualTo(
+      """
+      |package com.example
+      |
+      |import kotlin.text.isNullOrEmpty as textIsNullOrEmpty
+      |
+      |public fun main() {
+      |  val isFirstNull = null.textIsNullOrEmpty()
+      |  val isSecondNull = null.isNullOrEmpty()
+      |}
+      |
+      """.trimMargin(),
+    )
+  }
+
+  // not a good idea to do that, but still valid syntax
+  @Test fun extensionFunctionIsImportedEvenIfTheSameTypeIsAlreadyImported() {
+    val subpkgIsNullOrEmpty = ClassName(packageName = "com.example.subpkg", simpleNames = listOf("isNullOrEmpty"))
+    val kotlinIsNullOrEmpty = MemberName(packageName = "kotlin.text", simpleName = "isNullOrEmpty", isExtension = true)
+    val file = FileSpec.builder("com.example", "Test")
+      .addFunction(
+        FunSpec.builder("main")
+          .addStatement("val instance = %T()", subpkgIsNullOrEmpty)
+          .addStatement("val extensionFunctionResult = null.%M()", kotlinIsNullOrEmpty)
+          .build(),
+      )
+      .build()
+    assertThat(file.toString()).isEqualTo(
+      """
+      |package com.example
+      |
+      |import com.example.subpkg.isNullOrEmpty
+      |import kotlin.text.isNullOrEmpty
+      |
+      |public fun main() {
+      |  val instance = isNullOrEmpty()
+      |  val extensionFunctionResult = null.isNullOrEmpty()
+      |}
+      |
+      """.trimMargin(),
+    )
+  }
+
   // https://github.com/square/kotlinpoet/issues/1563
   @Test fun nestedClassesWithConflictingAutoGeneratedImports() {
     val source = FileSpec.builder("com.squareup.tacos", "Taco")