Update ChangeType's fix to not limit to Kotlin but more generic

openrewrite · Sep 23, 2023 · 4548736 · 4548736
1 parent 80712b6
commit 4548736
Showing 1 changed file with 14 additions and 32 deletions.
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/ChangeType.java b/rewrite-java/src/main/java/org/openrewrite/java/ChangeType.java
@@ -89,7 +89,6 @@ private static class ChangeTypeVisitor extends JavaVisitor<ExecutionContext> {
         @Nullable
         private J.Identifier importAlias;
         private boolean hasImportWithoutAlias;
-        private boolean isKotlinSource;
 
         @Nullable
         private final Boolean ignoreDefinition;
@@ -103,15 +102,11 @@ private ChangeTypeVisitor(String oldFullyQualifiedTypeName, String newFullyQuali
             this.ignoreDefinition = ignoreDefinition;
             importAlias = null;
             hasImportWithoutAlias = false;
-            isKotlinSource = false;
         }
 
         @Override
         public J visit(@Nullable Tree tree, ExecutionContext ctx) {
             if (tree instanceof JavaSourceFile) {
-                if (tree.getClass().getName().equals("org.openrewrite.kotlin.tree.K$CompilationUnit")) {
-                    isKotlinSource = true;
-                }
                 JavaSourceFile cu = (JavaSourceFile) requireNonNull(tree);
                 if (!Boolean.TRUE.equals(ignoreDefinition)) {
                     JavaType.FullyQualified fq = TypeUtils.asFullyQualified(targetType);
@@ -136,16 +131,14 @@ public J visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ct
 
         @Override
         public J visitImport(J.Import import_, ExecutionContext ctx) {
-            if (isKotlinSource) {
-                // Collect kotlin alias information here
-                // If there is an existing import with an alias, we need to add a target import with an alias accordingly.
-                // If there is an existing import without an alias, we need to add a target import with an alias accordingly.
-                if (hasSameFQN(import_, originalType)) {
-                    if (import_.getAlias() != null) {
-                        importAlias = import_.getAlias();
-                    } else {
-                        hasImportWithoutAlias = true;
-                    }
+            // Collect alias import information here
+            // If there is an existing import with an alias, we need to add a target import with an alias accordingly.
+            // If there is an existing import without an alias, we need to add a target import with an alias accordingly.
+            if (hasSameFQN(import_, originalType)) {
+                if (import_.getAlias() != null) {
+                    importAlias = import_.getAlias();
+                } else {
+                    hasImportWithoutAlias = true;
                 }
             }
 
@@ -160,7 +153,7 @@ public J visitImport(J.Import import_, ExecutionContext ctx) {
             return updateType(javaType);
         }
 
-        private void maybeAddKotlinImport(JavaType.FullyQualified owningClass) {
+        private void addImport(JavaType.FullyQualified owningClass) {
             if (importAlias != null) {
                 maybeAddImport(owningClass.getPackageName(), owningClass.getClassName(), null, importAlias.getSimpleName(), true);
             }
@@ -209,18 +202,10 @@ private void maybeAddKotlinImport(JavaType.FullyQualified owningClass) {
                     JavaType.FullyQualified owningClass = fullyQualifiedTarget.getOwningClass();
                     if (!(owningClass != null && topLevelClassnames.contains(getTopLevelClassName(fullyQualifiedTarget).getFullyQualifiedName()))) {
                         if (owningClass != null && !"java.lang".equals(fullyQualifiedTarget.getPackageName())) {
-                            if (isKotlinSource) {
-                                maybeAddKotlinImport(owningClass);
-                            } else {
-                                maybeAddImport(owningClass.getPackageName(), owningClass.getClassName(), null, null, true);
-                            }
+                            addImport(owningClass);
                         }
                         if (!"java.lang".equals(fullyQualifiedTarget.getPackageName())) {
-                            if (isKotlinSource) {
-                                maybeAddKotlinImport(fullyQualifiedTarget);
-                            } else {
-                                maybeAddImport(fullyQualifiedTarget.getPackageName(), fullyQualifiedTarget.getClassName(), null, null, true);
-                            }
+                            addImport(fullyQualifiedTarget);
                         }
                     }
                 }
@@ -322,11 +307,8 @@ public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx)
                                 method.getSimpleName().equals(anImport.getQualid().getSimpleName())) {
                                 JavaType.FullyQualified targetFqn = (JavaType.FullyQualified) targetType;
 
-                                if (isKotlinSource) {
-                                    maybeAddKotlinImport(targetFqn);
-                                } else {
-                                    maybeAddImport((targetFqn).getFullyQualifiedName(), method.getName().getSimpleName());
-                                }
+                                addImport(targetFqn);
+                                maybeAddImport((targetFqn).getFullyQualifiedName(), method.getName().getSimpleName());
                                 break;
                             }
                         }
@@ -639,7 +621,7 @@ public static JavaType.FullyQualified getTopLevelClassName(JavaType.FullyQualifi
         return getTopLevelClassName(classType.getOwningClass());
     }
 
-    public static boolean hasSameFQN(J.Import import_, JavaType targetType) {
+    private static boolean hasSameFQN(J.Import import_, JavaType targetType) {
         JavaType.FullyQualified type = TypeUtils.asFullyQualified(targetType);
         String fqn = type != null ? type.getFullyQualifiedName() : null;