Python: route imports through a PythonImportService so Java ChangeType stops corrupting files

rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java

@@ -42,6 +42,7 @@
 import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -75,7 +76,8 @@ public class RewriteRpc {
     /**
      * Keeps track of the local and remote state of objects that are used in
      * visits and other operations for which incremental state sharing is useful
-     * between two processes.
+     * between two processes. Note these do not need to be ConcurrentHashMap as
+     * each RewriteRpc instance is held in its own ThreadLocal in RewriteRpcProcessManager
      */
     @VisibleForTesting
     final Map<String, Object> remoteObjects = new HashMap<>();
@@ -372,6 +374,15 @@ public void evict(String id, int localRefsCheckpoint, int remoteRefsCheckpoint)
     }
 
     public <P> @Nullable Tree visit(Tree tree, String visitorName, P p, @Nullable Cursor cursor) {
+        return visit(tree, visitorName, null, p, cursor);
+    }
+
+    /**
+     * Run a remote visitor that takes constructor arguments (e.g. the peer's own {@code AddImport}),
+     * or {@code null} {@code visitorOptions} when it takes none.
+     */
+    public <P> @Nullable Tree visit(Tree tree, String visitorName, @Nullable Map<String, Object> visitorOptions,
+                                    P p, @Nullable Cursor cursor) {
         ensureDataTableStoreSent();
         // Set the local state of this tree, so that when the remote asks for it, we know what to send.
         localObjects.put(tree.getId().toString(), tree);
@@ -381,7 +392,7 @@ public void evict(String id, int localRefsCheckpoint, int remoteRefsCheckpoint)
 
         String sourceFileType = DynamicDispatchRpcCodec.canonicalSourceFileType(
                 (tree instanceof SourceFile ? tree : requireNonNull(cursor).firstEnclosingOrThrow(SourceFile.class)).getClass());
-        Supplier<VisitResponse> doSend = () -> send("Visit", new Visit(visitorName, sourceFileType, null,
+        Supplier<VisitResponse> doSend = () -> send("Visit", new Visit(visitorName, sourceFileType, visitorOptions,
                 tree.getId().toString(), pId, cursorIds), VisitResponse.class);
         VisitResponse response = p instanceof ExecutionContext
                 ? RewriteRpcExecutionContextView.view((ExecutionContext) p).withInFlightSlot(doSend)

rewrite-java/src/main/java/org/openrewrite/java/ChangeType.java

@@ -23,6 +23,7 @@
 import org.openrewrite.java.internal.PackageNameUtils;
 import org.openrewrite.java.marker.JavaSourceSet;
 import org.openrewrite.java.search.UsesType;
+import org.openrewrite.java.service.ImportService;
 import org.openrewrite.java.tree.*;
 import org.openrewrite.marker.Markers;
 import org.openrewrite.marker.SearchResult;
@@ -280,6 +281,11 @@ private void addImport(JavaType.FullyQualified owningClass) {
                     }
                 }
 
+                if (targetType instanceof JavaType.FullyQualified && service(ImportService.class).usesStatementBasedImports()) {
+                    // Queued after the addition so the import service sees the name already bound by
+                    // the new import, which is what makes the old one safe to retire.
+                    maybeRemoveImport(originalType.getFullyQualifiedName());
+                }
                 j = sf.withImports(ListUtils.map(sf.getImports(), i -> {
                     Cursor cursor = getCursor();
                     setCursor(new Cursor(cursor, i));

rewrite-java/src/main/java/org/openrewrite/java/JavaVisitor.java

@@ -199,7 +199,7 @@ public void maybeRemoveImport(JavaType.@Nullable FullyQualified clazz) {
     }
 
     public void maybeRemoveImport(String fullyQualifiedName) {
-        RemoveImport<P> op = new RemoveImport<>(fullyQualifiedName);
+        JavaVisitor<P> op = service(ImportService.class).removeImportVisitor(fullyQualifiedName);
         if (!getAfterVisit().contains(op)) {
             doAfterVisit(op);
         }

rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java

@@ -20,6 +20,7 @@
 import org.openrewrite.Incubating;
 import org.openrewrite.java.AddImport;
 import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.RemoveImport;
 import org.openrewrite.java.ShortenFullyQualifiedTypeReferences;
 import org.openrewrite.java.tree.J;
 
@@ -34,6 +35,20 @@ public <P> JavaVisitor<P> addImportVisitor(@Nullable String packageName,
         return new AddImport<>(packageName, typeName, member,  alias, onlyIfReferenced);
     }
 
+    public <P> JavaVisitor<P> removeImportVisitor(String fullyQualifiedName) {
+        return new RemoveImport<>(fullyQualifiedName);
+    }
+
+    /**
+     * Whether the language keeps its imports among the statements rather than in
+     * {@link org.openrewrite.java.tree.JavaSourceFile#getImports()}. Such a language exposes no
+     * {@link J.Import}, so callers that would otherwise walk that list have to route the edit
+     * through this service instead.
+     */
+    public boolean usesStatementBasedImports() {
+        return false;
+    }
+
     public <J2 extends J> JavaVisitor<ExecutionContext> shortenAllFullyQualifiedTypeReferences() {
         return new ShortenFullyQualifiedTypeReferences().getVisitor();
     }

rewrite-python/rewrite/src/rewrite/java/extensions.py

@@ -2,7 +2,7 @@
 
 from rewrite import Cursor
 from rewrite.tree import Tree
-from rewrite.utils import list_map
+from rewrite.utils import list_map, replace_if_changed
 from .support_types import J, JRightPadded, JLeftPadded, JContainer, Space, P, T, J2
 
 if TYPE_CHECKING:
@@ -68,4 +68,4 @@ def visit_space(v: 'JavaVisitor', space: Optional[Space], p: P) -> Space:
 
 def with_name(method: 'MethodInvocation', name: 'Identifier') -> 'MethodInvocation':
     # FIXME add type attribution logic
-    return method if name is method.name else replace(method, _name=name)
+    return replace_if_changed(method, _name=name)