fix: preserve comments and blank lines when removing unused imports

When removing an unused import, libcst's RemoveFromParent() also removed
comment lines and blank lines above the import. This change uses a two-phase
approach: mark imports for removal, then in leave_SimpleStatementLine extract
comment-bearing leading lines and transfer them to the next sibling statement.

A stack-based push/pop mechanism isolates pending lines across compound
statement boundaries (class, function, if, for, etc.) so nested statements
don't consume comments from the outer scope.

Closes #100

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

Author: hakancelikdev

src/unimport/refactor.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import cast
+from typing import Union, cast
 
 import libcst as cst
 import libcst.matchers as m
@@ -14,14 +14,17 @@
 
 
 class _RemoveUnusedImportTransformer(cst.CSTTransformer):
-    __slots__ = ("unused_imports",)
+    __slots__ = ("unused_imports", "_nodes_to_remove", "_pending_lines", "_pending_lines_stack")
 
     METADATA_DEPENDENCIES = (PositionProvider,)
 
     def __init__(self, unused_imports: list[Import | ImportFrom]) -> None:
         super().__init__()
 
         self.unused_imports = unused_imports
+        self._nodes_to_remove: set[int] = set()
+        self._pending_lines: list[cst.EmptyLine] = []
+        self._pending_lines_stack: list[list[cst.EmptyLine]] = []
 
     @staticmethod
     def get_import_name_from_attr(attr_node: cst.Attribute) -> str:
@@ -61,7 +64,7 @@ def get_rpar(rpar: cst.RightParen | None, location: CodeRange) -> cst.RightParen
 
     def leave_import_alike(
         self, original_node: T.CSTImportT, updated_node: T.CSTImportT
-    ) -> cst.RemovalSentinel | T.CSTImportT:
+    ) -> T.CSTImportT:
         names_to_keep = []
         names = cast(Sequence[cst.ImportAlias], updated_node.names)
         # already handled by leave_ImportFrom
@@ -78,7 +81,8 @@ def leave_import_alike(
             if self.is_import_used(import_name, column + 1, self.get_location(original_node)):
                 names_to_keep.append(import_alias)
         if not names_to_keep:
-            return cst.RemoveFromParent()
+            self._nodes_to_remove.add(id(original_node))
+            return updated_node
         elif len(names) == len(names_to_keep):
             return updated_node
         else:
@@ -91,19 +95,17 @@ def leave_import_alike(
             return cast(T.CSTImportT, updated_node)
 
     @staticmethod
-    def leave_StarImport(updated_node: cst.ImportFrom, imp: ImportFrom) -> cst.ImportFrom | cst.RemovalSentinel:
+    def leave_StarImport(updated_node: cst.ImportFrom, imp: ImportFrom) -> tuple[cst.ImportFrom, bool]:
         if imp.suggestions:
             names_to_suggestions = [cst.ImportAlias(cst.Name(module)) for module in imp.suggestions]
-            return updated_node.with_changes(names=names_to_suggestions)
+            return updated_node.with_changes(names=names_to_suggestions), False
         else:
-            return cst.RemoveFromParent()
+            return updated_node, True
 
-    def leave_Import(self, original_node: cst.Import, updated_node: cst.Import) -> cst.RemovalSentinel | cst.Import:
+    def leave_Import(self, original_node: cst.Import, updated_node: cst.Import) -> cst.Import:
         return self.leave_import_alike(original_node, updated_node)
 
-    def leave_ImportFrom(
-        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
-    ) -> cst.RemovalSentinel | cst.ImportFrom:
+    def leave_ImportFrom(self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom) -> cst.ImportFrom:
         if isinstance(updated_node.names, cst.ImportStar):
 
             def get_star_imp() -> ImportFrom | None:
@@ -120,12 +122,149 @@ def get_star_imp() -> ImportFrom | None:
 
             imp = get_star_imp()
             if imp:
-                return self.leave_StarImport(updated_node, imp)
+                result, should_remove = self.leave_StarImport(updated_node, imp)
+                if should_remove:
+                    self._nodes_to_remove.add(id(original_node))
+                return result
             else:
                 return original_node
 
         return self.leave_import_alike(original_node, updated_node)
 
+    def leave_SimpleStatementLine(
+        self,
+        original_node: cst.SimpleStatementLine,
+        updated_node: cst.SimpleStatementLine,
+    ) -> Union[cst.SimpleStatementLine, cst.RemovalSentinel]:
+        # Check if any child import node was marked for removal
+        should_remove = False
+        for stmt in original_node.body:
+            if id(stmt) in self._nodes_to_remove:
+                should_remove = True
+                break
+
+        if should_remove:
+            # Extract comment-bearing lines (and blank lines that precede them)
+            # from leading_lines and stash them
+            lines = list(updated_node.leading_lines)
+            preserved: list[cst.EmptyLine] = []
+            for i, line in enumerate(lines):
+                if isinstance(line, cst.EmptyLine) and line.comment is not None:
+                    # Also include blank lines immediately before this comment
+                    j = i - 1
+                    blank_prefix: list[cst.EmptyLine] = []
+                    while j >= 0 and isinstance(lines[j], cst.EmptyLine) and lines[j].comment is None:
+                        blank_prefix.append(lines[j])
+                        j -= 1
+                    blank_prefix.reverse()
+                    preserved.extend(blank_prefix)
+                    preserved.append(line)
+            self._pending_lines.extend(preserved)
+            return cst.RemoveFromParent()
+
+        # If there are pending comment lines, prepend them to this statement
+        if self._pending_lines:
+            new_leading = list(self._pending_lines) + list(updated_node.leading_lines)
+            self._pending_lines.clear()
+            return updated_node.with_changes(leading_lines=new_leading)
+
+        return updated_node
+
+    # -- Compound statement scope isolation --
+    # Push/pop pending lines so that nested statements don't consume
+    # pending lines from the outer scope.
+
+    def _push_pending(self) -> None:
+        self._pending_lines_stack.append(self._pending_lines)
+        self._pending_lines = []
+
+    def _pop_and_apply(self, updated_node: cst.BaseCompoundStatement) -> cst.BaseCompoundStatement:
+        self._pending_lines = self._pending_lines_stack.pop()
+        if self._pending_lines:
+            new_leading = list(self._pending_lines) + list(updated_node.leading_lines)
+            self._pending_lines.clear()
+            return updated_node.with_changes(leading_lines=new_leading)
+        return updated_node
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_FunctionDef(
+        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_If(self, node: cst.If) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_If(
+        self, original_node: cst.If, updated_node: cst.If
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_For(self, node: cst.For) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_For(
+        self, original_node: cst.For, updated_node: cst.For
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_While(self, node: cst.While) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_While(
+        self, original_node: cst.While, updated_node: cst.While
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_Try(self, node: cst.Try) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_Try(
+        self, original_node: cst.Try, updated_node: cst.Try
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_TryStar(self, node: cst.TryStar) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_TryStar(
+        self, original_node: cst.TryStar, updated_node: cst.TryStar
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def visit_With(self, node: cst.With) -> bool:
+        self._push_pending()
+        return True
+
+    def leave_With(
+        self, original_node: cst.With, updated_node: cst.With
+    ) -> Union[cst.BaseStatement, cst.RemovalSentinel]:
+        return self._pop_and_apply(updated_node)
+
+    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
+        if self._pending_lines:
+            new_footer = list(updated_node.footer) + list(self._pending_lines)
+            self._pending_lines.clear()
+            return updated_node.with_changes(footer=new_footer)
+        return updated_node
+
 
 def refactor_string(source: str, unused_imports: list[Import | ImportFrom]) -> str:
     if unused_imports:

tests/cases/analyzer/unused/comment_above_unused_before_class.py

@@ -0,0 +1,18 @@
+from typing import Union
+
+from unimport.statement import Import, ImportFrom, Name
+
+__all__ = ["NAMES", "IMPORTS", "UNUSED_IMPORTS"]
+
+
+NAMES: list[Name] = [
+    Name(lineno=9, name="print", is_all=False),
+    Name(lineno=9, name="os", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=1, column=1, name="os", package="os"),
+    Import(lineno=4, column=1, name="sys", package="sys"),
+]
+UNUSED_IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=4, column=1, name="sys", package="sys"),
+]

tests/cases/analyzer/unused/comment_above_unused_end_of_file.py

@@ -0,0 +1,18 @@
+from typing import Union
+
+from unimport.statement import Import, ImportFrom, Name
+
+__all__ = ["NAMES", "IMPORTS", "UNUSED_IMPORTS"]
+
+
+NAMES: list[Name] = [
+    Name(lineno=3, name="print", is_all=False),
+    Name(lineno=3, name="os", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=1, column=1, name="os", package="os"),
+    Import(lineno=6, column=1, name="sys", package="sys"),
+]
+UNUSED_IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=6, column=1, name="sys", package="sys"),
+]

tests/cases/analyzer/unused/comment_above_unused_import.py

@@ -0,0 +1,32 @@
+from typing import Union
+
+from unimport.statement import Import, ImportFrom, Name
+
+__all__ = ["NAMES", "IMPORTS", "UNUSED_IMPORTS"]
+
+
+NAMES: list[Name] = [
+    Name(lineno=6, name="print", is_all=False),
+    Name(lineno=6, name="os", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=1, column=1, name="os", package="os"),
+    ImportFrom(
+        lineno=4,
+        column=1,
+        name="compile_command",
+        package="codeop",
+        star=False,
+        suggestions=[],
+    ),
+]
+UNUSED_IMPORTS: list[Union[Import, ImportFrom]] = [
+    ImportFrom(
+        lineno=4,
+        column=1,
+        name="compile_command",
+        package="codeop",
+        star=False,
+        suggestions=[],
+    )
+]

tests/cases/analyzer/unused/comment_above_unused_import_with_siblings.py

@@ -0,0 +1,58 @@
+from typing import Union
+
+from unimport.statement import Import, ImportFrom, Name
+
+__all__ = ["NAMES", "IMPORTS", "UNUSED_IMPORTS"]
+
+
+NAMES: list[Name] = [
+    Name(lineno=11, name="os", is_all=False),
+    Name(lineno=11, name="Union", is_all=False),
+    Name(lineno=11, name="driver", is_all=False),
+    Name(lineno=11, name="Grammar", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=2, column=1, name="os", package="os"),
+    ImportFrom(
+        lineno=4,
+        column=1,
+        name="Union",
+        package="typing",
+        star=False,
+        suggestions=[],
+    ),
+    ImportFrom(
+        lineno=7,
+        column=1,
+        name="token",
+        package=".pgen2",
+        star=False,
+        suggestions=[],
+    ),
+    ImportFrom(
+        lineno=8,
+        column=1,
+        name="driver",
+        package=".pgen2",
+        star=False,
+        suggestions=[],
+    ),
+    ImportFrom(
+        lineno=10,
+        column=1,
+        name="Grammar",
+        package=".pgen2.grammar",
+        star=False,
+        suggestions=[],
+    ),
+]
+UNUSED_IMPORTS: list[Union[Import, ImportFrom]] = [
+    ImportFrom(
+        lineno=7,
+        column=1,
+        name="token",
+        package=".pgen2",
+        star=False,
+        suggestions=[],
+    )
+]

tests/cases/analyzer/unused/comment_unused_comment_used.py

@@ -0,0 +1,28 @@
+from typing import Union
+
+from unimport.statement import Import, ImportFrom, Name
+
+__all__ = ["NAMES", "IMPORTS", "UNUSED_IMPORTS"]
+
+
+NAMES: list[Name] = [
+    Name(lineno=9, name="print", is_all=False),
+    Name(lineno=9, name="os", is_all=False),
+    Name(lineno=10, name="print", is_all=False),
+    Name(lineno=10, name="OrderedDict", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=1, column=1, name="os", package="os"),
+    Import(lineno=4, column=1, name="sys", package="sys"),
+    ImportFrom(
+        lineno=7,
+        column=1,
+        name="OrderedDict",
+        package="collections",
+        star=False,
+        suggestions=[],
+    ),
+]
+UNUSED_IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=4, column=1, name="sys", package="sys"),
+]