From 01c700e4c92b56c683aa9bbbdcf96ab1318bf5b5 Mon Sep 17 00:00:00 2001
From: Dimitris Koutsogiorgas <dimitris@squareup.com>
Date: Mon, 6 Nov 2023 11:28:23 +0200
Subject: [PATCH] Fix imports when extending a type with a different return
 type

---
 .../java/io/outfoxx/swiftpoet/CodeWriter.kt   |  4 +-
 .../outfoxx/swiftpoet/test/FileSpecTests.kt   | 45 ++++++++++++++++++-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt b/src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt
index afa2d5da6..c79d4c712 100644
--- a/src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt
+++ b/src/main/java/io/outfoxx/swiftpoet/CodeWriter.kt
@@ -329,7 +329,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
index 3feba66ac..cd581aae7 100644
--- 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 {
               }
             
             }