Merge pull request #90 from outfoxx/extension_import_bug

Fix emitting imports using ExternalTypeSpecs
outfoxx · Dec 14, 2023 · 29299c2 · 29299c2
2 parents dc661a0 + 01c700e
commit 29299c2
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt b/src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt
@@ -336,7 +336,9 @@ internal class CodeWriter constructor(
     for (i in typeSpecStack.indices.reversed()) {
       val typeSpec = typeSpecStack[i]
       if (typeSpec is ExternalTypeSpec) {
-        return stackTypeName(i, simpleName)
+        if (typeSpec.name == simpleName) {
+          return stackTypeName(i, simpleName)
+        }
       }
       for (visibleChild in typeSpec.typeSpecs) {
         if (visibleChild.name == simpleName) {

diff --git a/src/test/java/io/outfoxx/swiftpoet/test/FileSpecTests.kt b/src/test/java/io/outfoxx/swiftpoet/test/FileSpecTests.kt
@@ -79,9 +79,50 @@ class FileSpecTests {
     )
   }
 
+  @Test
+  @DisplayName("Generates correct imports when extending different types")
+  fun testImportsForDifferentExtensionTypes() {
+    val parentElement = typeName("Foundation.Data")
+    val obsElement = typeName("RxSwift.Observable.Element")
+
+    val extension =
+      ExtensionSpec.builder(parentElement)
+        .addFunction(
+          FunctionSpec.builder("test")
+            .returns(obsElement)
+            .build()
+        )
+        .build()
+
+    val testFile = FileSpec.builder("Test", "Test")
+      .addExtension(extension)
+      .build()
+
+    val out = StringWriter()
+    testFile.writeTo(out)
+
+    assertThat(
+      out.toString(),
+      equalTo(
+        """
+            import Foundation
+            import RxSwift
+
+            extension Data {
+
+              func test() -> Observable.Element {
+              }
+
+            }
+
+        """.trimIndent()
+      )
+    )
+  }
+
   @Test
   @DisplayName("Generates correct imports for extension types")
-  fun testImportsAndShortensExtensionTypes() {
+  fun testImportsForSameExtensionTypes() {
 
     val obsElement = typeName("RxSwift.Observable.Element")
 
@@ -109,7 +150,7 @@ class FileSpecTests {
 
             extension Observable {
             
-              func test() -> Element {
+              func test() -> RxSwift.Observable.Element {
               }
             
             }