Fix import after use in function incorrectly removed (#178) (#321)

When a module-level import appears textually after a function that uses
it, unimport incorrectly removed it. Function/async function bodies are
deferred (only execute when called), so imports defined later at module
scope are available at call time. Class bodies execute immediately, so
the lineno check still applies for them.

Added Name._is_deferred_usage() that walks from the Name's scope toward
the Import's scope — if any scope on the path is a FunctionDef or
AsyncFunctionDef, the usage is deferred and the lineno check is relaxed.

Edge case tests cover: async functions, nested functions, class bodies
(not deferred), methods inside classes, from-imports, and module-level
sequential usage.

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

Author: hakancelikdev

src/unimport/statement.py

@@ -140,15 +140,26 @@ class Name:
     def is_attribute(self):
         return "." in self.name
 
+    def _is_deferred_usage(self, imp: Import | ImportFrom) -> bool:
+        imp_scope = imp.scope
+        scope = self.scope
+        while scope is not None and scope != imp_scope:
+            if isinstance(scope.node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                return True
+            scope = scope.parent
+        return False
+
     def match_2(self, imp: Import | ImportFrom) -> bool:
         if self.is_all:
             is_match = self.name == imp.name
         elif self.is_attribute:
-            is_match = imp.lineno < self.lineno and (
+            is_match = (imp.lineno < self.lineno or self._is_deferred_usage(imp)) and (
                 ".".join(self.name.split(".")[: len(imp)]) == imp.name or imp.is_match_sub_packages(self.name)
             )
         else:
-            is_match = (imp.lineno < self.lineno) and (self.name == imp.name or imp.is_match_sub_packages(self.name))
+            is_match = (imp.lineno < self.lineno or self._is_deferred_usage(imp)) and (
+                self.name == imp.name or imp.is_match_sub_packages(self.name)
+            )
 
         return is_match
 

tests/cases/analyzer/unused/from_import_after_use_in_function.py

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

tests/cases/analyzer/unused/import_after_use_at_module_level.py

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

tests/cases/analyzer/unused/import_after_use_in_async_function.py

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

tests/cases/analyzer/unused/import_after_use_in_class.py

@@ -0,0 +1,17 @@
+from typing import Union
+
+from unimport.statement import Import, ImportFrom, Name
+
+__all__ = ["NAMES", "IMPORTS", "UNUSED_IMPORTS"]
+
+
+NAMES: list[Name] = [
+    Name(lineno=2, name="x", is_all=False),
+    Name(lineno=2, name="sys.platform", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    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/import_after_use_in_function.py

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

tests/cases/analyzer/unused/import_after_use_in_method.py

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

tests/cases/analyzer/unused/import_after_use_in_nested_function.py

@@ -0,0 +1,16 @@
+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.path.join", is_all=False),
+    Name(lineno=3, name="os.path", is_all=False),
+]
+IMPORTS: list[Union[Import, ImportFrom]] = [
+    Import(lineno=5, column=1, name="os", package="os"),
+]
+UNUSED_IMPORTS: list[Union[Import, ImportFrom]] = []

tests/cases/refactor/unused/from_import_after_use_in_function.py

@@ -0,0 +1,4 @@
+def foo():
+    return path.join("a", "b")
+
+from os import path

tests/cases/refactor/unused/import_after_use_at_module_level.py

@@ -0,0 +1 @@
+print(sys.path)