Fix import parsing

openrewrite · Nov 8, 2023 · 819c8ec · 819c8ec
1 parent 13f2948
commit 819c8ec
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -2361,10 +2361,16 @@ public J visitImportDirective(KtImportDirective importDirective, ExecutionContex
         ASTNode node = importDirective.getNode().findChildByType(KtTokens.IMPORT_KEYWORD);
         LeafPsiElement importPsi = (LeafPsiElement) node;
 
-        String text = nodeRangeText(
-                findFirstNonSpaceNextSibling(importPsi).getNode(),
-                importDirective.isAllUnder() ? importDirective.getNode().findChildByType(KtTokens.MUL)
-                        : importDirective.getNode().findChildByType(KtNodeTypes.DOT_QUALIFIED_EXPRESSION));
+        PsiElement first = findFirstNonSpaceNextSibling(importPsi);
+        ASTNode last = importDirective.isAllUnder() ? importDirective.getNode().findChildByType(KtTokens.MUL)
+                : importDirective.getNode().findChildByType(KtNodeTypes.DOT_QUALIFIED_EXPRESSION);
+        if (last == null) {
+            PsiElement lastPsi = findLastNotSpaceChild(importDirective);
+            if (lastPsi != null) {
+                last = lastPsi.getNode();
+            }
+        }
+        String text = nodeRangeText(first != null ? first.getNode() : null, last);
         J reference = TypeTree.build(text); // FIXME: this creates a shallow class for a resolvable type.
         reference = reference.withPrefix(suffix(importPsi));
 

diff --git a/src/test/java/org/openrewrite/kotlin/tree/ImportTest.java b/src/test/java/org/openrewrite/kotlin/tree/ImportTest.java
@@ -96,7 +96,7 @@ class A {
     void methodName() {
         rewriteRun(
           kotlin("fun <T : Any> Class<T>.createInstance() {}"),
-          kotlin("import   createInstance")
+          kotlin("import   createInstance /*C1*/")
         );
     }