Remove unused imports and remove autoflake dependency

.github/workflows/validate-code.yml

@@ -37,9 +37,6 @@ jobs:
           chmod +x ./validate.sh
           ./validate.sh
 
-      - name: Install Dependencies
-        run: pip install -r requirements.txt
-
       - name: Unit Tests
         run: |
           chmod +x ./test.sh

README.md

@@ -1,6 +1,6 @@
 # Personal Python AST Optimizer
 
-This is a project related to minifying or otherwise statically optimizing Python code. It can parse python to output a much smaller output with optionally excluded sections of code. It also includes wrappers for calling autoflake.
+This is a project related to minifying or otherwise statically optimizing Python code. It can parse python to output a much smaller output with optionally excluded sections of code.
 
 This project is an extension of the builtin ast module.
 

personal_python_ast_optimizer/parser/config.py

@@ -123,6 +123,7 @@ class OptimizationsConfig(_Config):
     __slots__ = (
         "vars_to_fold",
         "enums_to_fold",
+        "remove_unused_imports",
         "fold_constants",
         "assume_this_machine",
     )
@@ -132,6 +133,7 @@ def __init__(
         vars_to_fold: dict[str, int | str | bool] | None = None,
         enums_to_fold: Iterable[EnumType] | None = None,
         fold_constants: bool = True,
+        remove_unused_imports: bool = True,
         assume_this_machine: bool = False,
     ) -> None:
         self.vars_to_fold: dict[str, int | str | bool] = (
@@ -142,6 +144,7 @@ def __init__(
             if enums_to_fold is None
             else self._format_enums_to_fold_as_dict(enums_to_fold)
         )
+        self.remove_unused_imports: bool = remove_unused_imports
         self.assume_this_machine: bool = assume_this_machine
         self.fold_constants: bool = fold_constants
 

personal_python_ast_optimizer/parser/skipper.py

@@ -13,6 +13,7 @@
     machine_dependent_functions,
 )
 from personal_python_ast_optimizer.parser.utils import (
+    exclude_imports,
     first_occurrence_of_type,
     get_node_name,
     is_overload_function,
@@ -24,9 +25,10 @@
 )
 
 
-class AstNodeSkipper(ast.NodeTransformer):
+class AstNodeSkipper(ast.NodeVisitor):
 
     __slots__ = (
+        "_possibly_unused_imports",
         "_skippable_futures",
         "_within_class",
         "_within_function",
@@ -60,6 +62,8 @@ def __init__(self, config: SkipConfig) -> None:
         if self.token_types_config.skip_type_hints:
             self._skippable_futures.append("annotations")
 
+        self._possibly_unused_imports: set[str] = set()
+
     @staticmethod
     def _within_class_node(function):
         def wrapper(self: "AstNodeSkipper", *args, **kwargs) -> ast.AST | None:
@@ -100,9 +104,9 @@ def generic_visit(self, node) -> ast.AST:
                     new_values.append(value)
 
                 if (
-                    field == "body"
+                    not isinstance(node, ast.Module)
                     and not new_values
-                    and not isinstance(node, ast.Module)
+                    and field == "body"
                 ):
                     new_values = [ast.Pass()]
 
@@ -114,6 +118,7 @@ def generic_visit(self, node) -> ast.AST:
                     delattr(node, field)
                 else:
                     setattr(node, field, new_node)
+
         return node
 
     @staticmethod
@@ -145,13 +150,19 @@ def visit_Module(self, node: ast.Module) -> ast.AST:
         if not self._has_code_to_skip():
             return node
 
+        self._possibly_unused_imports = set()
+
         if self.token_types_config.skip_dangling_expressions:
             skip_dangling_expressions(node)
 
-        try:
-            return self.generic_visit(node)
-        finally:
-            self._warn_unused_skips()
+        module: ast.Module = self.generic_visit(node)  # type:ignore
+
+        if self._possibly_unused_imports:
+            import_skipper = ImportSkipper(self._possibly_unused_imports)
+            import_skipper.generic_visit(module)
+
+        self._warn_unused_skips()
+        return module
 
     @_within_class_node
     def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST | None:
@@ -232,6 +243,9 @@ def visit_Attribute(self, node: ast.Attribute) -> ast.AST | None:
         ):
             return self._get_enum_value_as_AST(node.value.attr, node.attr)
 
+        if node.attr in self._possibly_unused_imports:
+            self._possibly_unused_imports.remove(node.attr)
+
         return self.generic_visit(node)
 
     def _get_enum_value_as_AST(self, class_name: str, value_name: str) -> ast.Constant:
@@ -343,50 +357,44 @@ def visit_AugAssign(self, node: ast.AugAssign) -> ast.AST | None:
         return self.generic_visit(node)
 
     def visit_Import(self, node: ast.Import) -> ast.AST | None:
-        """Removes imports provided in config, deleting the whole
-        node if no imports are left"""
-        node.names = [
-            alias
-            for alias in node.names
-            if alias.name not in self.tokens_config.module_imports_to_skip
-        ]
+        exclude_imports(node, self.tokens_config.module_imports_to_skip)
 
         if not node.names:
             return None
 
+        if self.optimizations_config.remove_unused_imports:
+            self._possibly_unused_imports.update(n.asname or n.name for n in node.names)
+
         return self.generic_visit(node)
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST | None:
         normalized_module_name: str = node.module or ""
         if normalized_module_name in self.tokens_config.module_imports_to_skip:
             return None
 
-        node.names = [
-            alias
-            for alias in node.names
-            if alias.name not in self.tokens_config.from_imports_to_skip
-            and alias.name not in self.optimizations_config.vars_to_fold
-            and alias.name not in self.optimizations_config.enums_to_fold
-        ]
+        exclude_imports(node, self.tokens_config.from_imports_to_skip)
 
         if node.module == "__future__" and self._skippable_futures:
-            node.names = [
-                alias
-                for alias in node.names
-                if alias.name not in self._skippable_futures
-            ]
-
-        if not node.names:
-            return None
+            exclude_imports(node, self._skippable_futures)
+        elif (
+            node.module != "__future__"
+            and node.names
+            and self.optimizations_config.remove_unused_imports
+        ):
+            self._possibly_unused_imports.update(n.asname or n.name for n in node.names)
 
-        return self.generic_visit(node)
+        return self.generic_visit(node) if node.names else None
 
     def visit_Name(self, node: ast.Name) -> ast.AST:
         """Extends super's implementation by adding constant folding"""
         if node.id in self.optimizations_config.vars_to_fold:
             constant_value = self.optimizations_config.vars_to_fold[node.id]
+
             return ast.Constant(constant_value)
         else:
+            if node.id in self._possibly_unused_imports:
+                self._possibly_unused_imports.remove(node.id)
+
             return self.generic_visit(node)
 
     def visit_Dict(self, node: ast.Dict) -> ast.AST:
@@ -649,3 +657,45 @@ def _ast_constants_operation(
                 raise ValueError(f"Invalid operation: {operation.__class__.__name__}")
 
         return ast.Constant(result)
+
+
+class ImportSkipper(ast.NodeVisitor):
+
+    __slots__ = ("unused_imports",)
+
+    def __init__(self, unused_imports: set[str]) -> None:
+        self.unused_imports = unused_imports
+
+    def generic_visit(self, node):
+        for _, old_value in ast.iter_fields(node):
+            if isinstance(old_value, list):
+                new_values = []
+                ast_removed: bool = False
+                for value in old_value:
+                    if isinstance(value, ast.AST):
+                        value = self.visit(value)
+                        if value is None:
+                            ast_removed = True
+                            continue
+                        elif not isinstance(value, ast.AST):
+                            new_values.extend(value)
+                            continue
+
+                    new_values.append(value)
+
+                if not isinstance(node, ast.Module) and not new_values and ast_removed:
+                    new_values = [ast.Pass()]
+
+                old_value[:] = new_values
+
+        return node
+
+    def visit_Import(self, node: ast.Import) -> ast.Import | None:
+        exclude_imports(node, self.unused_imports)
+
+        return node if node.names else None
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.ImportFrom | None:
+        exclude_imports(node, self.unused_imports)
+
+        return node if node.names else None

personal_python_ast_optimizer/parser/utils.py

@@ -5,6 +5,12 @@
 from personal_python_ast_optimizer.parser.config import TokensToSkip
 
 
+def exclude_imports(node: ast.Import | ast.ImportFrom, exlcudes: Iterable[str]) -> None:
+    node.names = [
+        alias for alias in node.names if (alias.asname or alias.name) not in exlcudes
+    ]
+
+
 def get_node_name(node: ast.AST | None) -> str:
     """Gets id or attr which both can represent var names"""
     if isinstance(node, ast.Call):

pyproject.toml

@@ -5,8 +5,7 @@ authors         = [
 name            = "personal-python-ast-optimizer"
 readme          = "README.md"
 requires-python = ">=3.11"
-dynamic         = ["dependencies", "version"]
+dynamic         = ["version"]
 
 [tool.setuptools.dynamic]
-dependencies = {file = ["requirements.txt"]}
 version         = {file = "version.txt"}

tests/parser/test_enum_folding.py

@@ -38,9 +38,7 @@ class _SomeIntEnum(IntEnum):
 import somewhere
 
 print(somewhere.someModule._SomeStrEnum.C)""",
-        """
-import somewhere
-print('C')""".strip(),
+        "print('C')",
     ),
 ]
 

tests/parser/test_excludes.py

@@ -1,6 +1,9 @@
 import pytest
 
-from personal_python_ast_optimizer.parser.config import TokensConfig
+from personal_python_ast_optimizer.parser.config import (
+    OptimizationsConfig,
+    TokensConfig,
+)
 from tests.utils import BeforeAndAfter, run_minifier_and_assert_correct
 
 
@@ -160,12 +163,18 @@ def test_exclude_module_imports():
 import numpy
 from numpy._core import uint8
 from . import asdf
+import a
+import a as b
+import a as c
 """,
-        "",
+        "import a as c",
     )
     run_minifier_and_assert_correct(
         before_and_after,
-        tokens_config=TokensConfig(module_imports_to_skip={"numpy", "numpy._core", ""}),
+        tokens_config=TokensConfig(
+            module_imports_to_skip={"numpy", "numpy._core", "", "a", "b"}
+        ),
+        optimizations_config=OptimizationsConfig(remove_unused_imports=False),
     )
 
 
@@ -187,4 +196,5 @@ def test_exclude_real_case():
         tokens_config=TokensConfig(
             functions_to_skip={"getLogger"}, variables_to_skip={"TYPE_CHECKING"}
         ),
+        optimizations_config=OptimizationsConfig(remove_unused_imports=False),
     )

tests/parser/test_imports.py

@@ -1,6 +1,9 @@
 import pytest
 
-from personal_python_ast_optimizer.parser.config import TokenTypesConfig
+from personal_python_ast_optimizer.parser.config import (
+    OptimizationsConfig,
+    TokenTypesConfig,
+)
 from tests.utils import BeforeAndAfter, run_minifier_and_assert_correct
 
 _futures_imports: str = """
@@ -54,7 +57,10 @@ def i():
         """import test,test2
 def i():import a,d;from b import c,d as e;from .b import f;print();import e""",
     )
-    run_minifier_and_assert_correct(before_and_after)
+    run_minifier_and_assert_correct(
+        before_and_after,
+        optimizations_config=OptimizationsConfig(remove_unused_imports=False),
+    )
 
 
 def test_import_star():
@@ -63,4 +69,21 @@ def test_import_star():
         "from ctypes import *",
         "from ctypes import*",
     )
+    run_minifier_and_assert_correct(
+        before_and_after,
+        optimizations_config=OptimizationsConfig(remove_unused_imports=False),
+    )
+
+
+def test_remove_unused_imports():
+
+    before_and_after = BeforeAndAfter(
+        """
+if a == b:
+    import foo
+    import bar
+
+print(a)""",
+        "if a==b:pass\nprint(a)",
+    )
     run_minifier_and_assert_correct(before_and_after)

tests/parser/test_real.py

@@ -11,7 +11,7 @@ def test_image_viewer_constants():
 File with constants needed in multiple spots of the codebase
 \"\"\"
 
-from enum import IntEnum, StrEnum
+from enum import StrEnum
 
 
 class ImageFormats(StrEnum):
@@ -22,7 +22,7 @@ class ImageFormats(StrEnum):
     JPEG = "JPEG"
     PNG = "PNG"
     WEBP = "WebP\"""",
-        """from enum import IntEnum,StrEnum
+        """from enum import StrEnum
 class ImageFormats(StrEnum):DDS='DDS';GIF='GIF';JPEG='JPEG';PNG='PNG';WEBP='WebP'""",
     )
     run_minifier_and_assert_correct(before_and_after)