Python: changing how imports are parsed, prefer J.Import (#6730)

* Python: Use J.Import for single imports instead of MultiImport

Single `import foo` statements no longer get unnecessarily wrapped in
a `MultiImport` node. MultiImport is now only used when there are
multiple names (`import foo, bar`). Updated parser, printer, and
import recipes (add, remove, change) accordingly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Added assertion

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: greg-at-moderne

rewrite-python/rewrite/src/rewrite/python/_parser_visitor.py

@@ -789,7 +789,11 @@ def visit_Try(self, node):
         )
 
     def visit_Import(self, node):
-        # TODO only use `MultiImport` when necessary (requires corresponding changes to printer)
+        if len(node.names) == 1:
+            prefix = self.__source_before('import')
+            imp = self.__convert(node.names[0])
+            return imp.replace(prefix=prefix, qualid=imp.qualid.replace(prefix=imp.prefix))
+
         prefix = self.__source_before('import')
         names_prefix = self.__whitespace()
         return py.MultiImport(

rewrite-python/rewrite/src/rewrite/python/add_import.py

@@ -131,7 +131,10 @@ def visit_compilation_unit(self, cu: CompilationUnit, p) -> J:
     def _import_exists(self, cu: CompilationUnit) -> bool:
         """Check if the import already exists."""
         for stmt in cu.statements:
-            if isinstance(stmt, MultiImport):
+            if isinstance(stmt, Import) and not isinstance(stmt, MultiImport):
+                if self.name is None and self._import_name_matches(stmt, self.module, self.alias):
+                    return True
+            elif isinstance(stmt, MultiImport):
                 if self._multi_import_matches(stmt):
                     return True
         return False

rewrite-python/rewrite/src/rewrite/python/printer.py

@@ -1475,7 +1475,14 @@ def visit_import(self, import_: 'j.Import', p: PrintOutputCapture) -> J:
 
         self._before_syntax(import_, p)
 
+        from rewrite.python import tree as _py
+        is_standalone = not self.get_cursor().first_enclosing(_py.MultiImport)
+        if is_standalone:
+            p.append("import")
+
         if isinstance(import_.qualid.target, j.Empty):
+            if is_standalone:
+                self._visit_space(import_.qualid.prefix, p)
             self.visit(import_.qualid.name, p)
         else:
             self.visit(import_.qualid, p)

rewrite-python/rewrite/src/rewrite/python/recipes/change_import.py

@@ -131,7 +131,13 @@ def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J:
                 self.old_alias = None
 
                 for stmt in cu.statements:
-                    if isinstance(stmt, MultiImport):
+                    if isinstance(stmt, Import) and not isinstance(stmt, MultiImport):
+                        alias = self._check_for_old_single_import(stmt)
+                        if alias is not None:
+                            self.has_old_import = True
+                            self.old_alias = alias if alias != "" else None
+                            break
+                    elif isinstance(stmt, MultiImport):
                         alias = self._check_for_old_import(stmt)
                         if alias is not None:
                             self.has_old_import = True
@@ -164,6 +170,16 @@ def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J:
 
                 return result
 
+            def visit_import(self, imp: Import, p: ExecutionContext) -> Optional[J]:
+                if not self.has_old_import or old_name:
+                    return imp
+                if self.cursor.first_enclosing(MultiImport):
+                    return imp
+                alias = self._check_for_old_single_import(imp)
+                if alias is None:
+                    return imp
+                return None
+
             def visit_multi_import(self, multi: MultiImport, p: ExecutionContext) -> Optional[J]:  # ty: ignore[invalid-method-override]
                 if not self.has_old_import:
                     return multi
@@ -180,6 +196,15 @@ def visit_multi_import(self, multi: MultiImport, p: ExecutionContext) -> Optiona
                     # import X - remove entire import
                     return self._remove_module_from_import(multi, old_module)
 
+            def _check_for_old_single_import(self, imp: Import) -> Optional[str]:
+                """Check if a standalone J.Import matches the old import."""
+                if old_name:
+                    return None
+                name = self._get_qualid_name(imp.qualid)
+                if name == old_module:
+                    return self._get_alias_name(imp) or ""
+                return None
+
             def _check_for_old_import(self, multi: MultiImport) -> Optional[str]:
                 """Check if this MultiImport matches the old import.
 

rewrite-python/rewrite/src/rewrite/python/remove_import.py

@@ -135,6 +135,20 @@ def _remove_import(self, cu: CompilationUnit) -> CompilationUnit:
 
         for padded in cu.padding.statements:
             stmt = padded.element
+            if isinstance(stmt, Import) and not isinstance(stmt, MultiImport):
+                result = self._process_single_import(stmt)
+                if result is None:
+                    removed_prefix = stmt.prefix
+                    changed = True
+                else:
+                    if removed_prefix is not None:
+                        padded = JRightPadded(
+                            padded.element.replace(prefix=removed_prefix),
+                            padded.after, padded.markers
+                        )
+                        removed_prefix = None
+                    new_padded_stmts.append(padded)
+                continue
             if isinstance(stmt, MultiImport):
                 result = self._process_multi_import(stmt)
                 if result is None:
@@ -169,6 +183,15 @@ def _remove_import(self, cu: CompilationUnit) -> CompilationUnit:
             return cu.padding.replace(_statements=new_padded_stmts)
         return cu
 
+    def _process_single_import(self, imp: Import) -> Optional[Import]:
+        """Process a standalone J.Import. Return None to remove, or the original."""
+        if self.name is not None:
+            return imp
+        name = self._get_qualid_name(imp.qualid)
+        if name == self.module:
+            return None
+        return imp
+
     def _process_multi_import(self, multi: MultiImport) -> Optional[MultiImport]:
         """Process a MultiImport and return None to remove, modified, or original."""
         if self.name is None:

rewrite-python/rewrite/tests/python/all/tree/import_test.py

@@ -1,10 +1,19 @@
+from rewrite.java import Import
+from rewrite.python import CompilationUnit
 from rewrite.test import RecipeSpec, python
 
 
+def _assert_single_j_import(cu: CompilationUnit) -> None:
+    assert len(cu.statements) == 1
+    stmt = cu.statements[0]
+    assert isinstance(stmt, Import), \
+        f"Single import should be a J.Import, got {type(stmt).__name__}"
+
+
 # noinspection PyUnresolvedReferences
 def test_simple():
     # language=python
-    RecipeSpec().rewrite_run(python("import io"))
+    RecipeSpec().rewrite_run(python("import io", after_recipe=_assert_single_j_import))
 
 
 # noinspection PyUnresolvedReferences
@@ -22,7 +31,7 @@ def test_unicode_char_normalization():
 # noinspection PyUnresolvedReferences
 def test_qualified():
     # language=python
-    RecipeSpec().rewrite_run(python("import xml.dom"))
+    RecipeSpec().rewrite_run(python("import xml.dom", after_recipe=_assert_single_j_import))
 
 
 # noinspection PyUnresolvedReferences