Add ImportService#shortenFullyQualifiedTypeReferencesIn() (#3687)

* Add `ImportService#shortenFullyQualifiedTypeReferencesIn()` Extends `ImportService` with `shortenFullyQualifiedTypeReferencesIn()` so that other languages extending `J` can provide a customized implementation. * Add some Javadoc
openrewrite · Nov 16, 2023 · efd3214 · efd3214
1 parent 9538741
commit efd3214
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
diff --git a/...java-test/src/test/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferencesTest.java b/...java-test/src/test/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferencesTest.java
@@ -18,6 +18,7 @@
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.ExecutionContext;
+import org.openrewrite.java.service.ImportService;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.JavaType;
 import org.openrewrite.test.RecipeSpec;
@@ -228,7 +229,7 @@ void visitSubtreeOnly() {
               @SuppressWarnings("DataFlowIssue")
               public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
                   if (method.getSimpleName().equals("m1")) {
-                      return (J.MethodDeclaration) new ShortenFullyQualifiedTypeReferences().getVisitor().visit(method, ctx, getCursor().getParent());
+                      return (J.MethodDeclaration) ShortenFullyQualifiedTypeReferences.modifyOnly(method).visit(method, ctx, getCursor().getParent());
                   }
                   return super.visitMethodDeclaration(method, ctx);
               }
@@ -429,7 +430,7 @@ void subtree() {
               @Override
               public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
                   if (method.getSimpleName().equals("m1")) {
-                      doAfterVisit(ShortenFullyQualifiedTypeReferences.modifyOnly(method));
+                      doAfterVisit(service(ImportService.class).shortenFullyQualifiedTypeReferencesIn(method));
                   }
                   return super.visitMethodDeclaration(method, ctx);
               }

diff --git a/rewrite-java/src/main/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferences.java b/rewrite-java/src/main/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferences.java
@@ -19,6 +19,7 @@
 import org.openrewrite.internal.lang.NonNull;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.internal.DefaultJavaTypeSignatureBuilder;
+import org.openrewrite.java.service.ImportService;
 import org.openrewrite.java.tree.*;
 
 import java.time.Duration;
@@ -47,11 +48,31 @@ public String getDescription() {
     }
 
     @Override
-    public TreeVisitor<?, ExecutionContext> getVisitor() {
-        return getVisitor(null);
+    public JavaVisitor<ExecutionContext> getVisitor() {
+        // This wrapper is necessary so that the "correct" implementation is used when this recipe is used declaratively
+        return new JavaVisitor<ExecutionContext>() {
+            @Override
+            public @Nullable J visit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof JavaSourceFile) {
+                    return ((JavaSourceFile) tree).service(ImportService.class).shortenFullyQualifiedTypeReferencesIn((J) tree).visit(tree, ctx);
+                }
+                return (J) tree;
+            }
+        };
     }
 
-    public static <J2 extends J> TreeVisitor<J, ExecutionContext> modifyOnly(J2 subtree) {
+    /**
+     * Returns a visitor which replaces all fully qualified references in the given subtree with simple names and adds 
+     * corresponding import statements.
+     * <p>
+     * For compatibility with other Java-based languages it is recommended to use this as a service via
+     * {@link ImportService#shortenFullyQualifiedTypeReferencesIn(J)}, as that will dispatch to the correct
+     * implementation for the language.
+     * 
+     * @see ImportService#shortenFullyQualifiedTypeReferencesIn(J)
+     * @see JavaVisitor#service(Class) 
+     */
+    public static <J2 extends J> JavaVisitor<ExecutionContext> modifyOnly(J2 subtree) {
         return getVisitor(subtree);
     }
 

diff --git a/rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java b/rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java
@@ -15,10 +15,13 @@
  */
 package org.openrewrite.java.service;
 
+import org.openrewrite.ExecutionContext;
 import org.openrewrite.Incubating;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.AddImport;
 import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.ShortenFullyQualifiedTypeReferences;
+import org.openrewrite.java.tree.J;
 
 @Incubating(since = "8.2.0")
 public class ImportService {
@@ -30,4 +33,12 @@ public <P> JavaVisitor<P> addImportVisitor(@Nullable String packageName,
                                                boolean onlyIfReferenced) {
         return new AddImport<>(packageName, typeName, member,  alias, onlyIfReferenced);
     }
+
+    public <J2 extends J> JavaVisitor<ExecutionContext> shortenAllFullyQualifiedTypeReferences() {
+        return new ShortenFullyQualifiedTypeReferences().getVisitor();
+    }
+
+    public <J2 extends J> JavaVisitor<ExecutionContext> shortenFullyQualifiedTypeReferencesIn(J2 subtree) {
+        return ShortenFullyQualifiedTypeReferences.modifyOnly(subtree);
+    }
 }