style.py: add import-check for missing & redundant imports (#47619)

3 files changed, 243 insertions, 22 deletions

diff --git a/lib/spack/spack/cmd/style.py b/lib/spack/spack/cmd/style.py
index d16ff34277..5925d1120e 100644
--- a/lib/spack/spack/cmd/style.py
+++ b/lib/spack/spack/cmd/style.py
@@ -3,10 +3,12 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 import argparse
+import ast
 import os
 import re
 import sys
 from itertools import zip_longest
+from typing import Dict, List, Optional
 
 import llnl.util.tty as tty
 import llnl.util.tty.color as color
@@ -14,7 +16,7 @@ from llnl.util.filesystem import working_dir
 
 import spack.paths
 import spack.util.git
-from spack.util.executable import which
+from spack.util.executable import Executable, which
 
 description = "runs source code style checks on spack"
 section = "developer"
@@ -36,10 +38,7 @@ exclude_directories = [os.path.relpath(spack.paths.external_path, spack.paths.pr
 #: double-check the results of other tools (if, e.g., --fix was provided)
 #: The list maps an executable name to a method to ensure the tool is
 #: bootstrapped or present in the environment.
-tool_names = ["isort", "black", "flake8", "mypy"]
-
-#: tools we run in spack style
-tools = {}
+tool_names = ["import-check", "isort", "black", "flake8", "mypy"]
 
 #: warnings to ignore in mypy
 mypy_ignores = [
@@ -61,14 +60,28 @@ def is_package(f):
 
 #: decorator for adding tools to the list
 class tool:
-    def __init__(self, name, required=False):
+    def __init__(self, name: str, required: bool = False, external: bool = True) -> None:
         self.name = name
+        self.external = external
         self.required = required
 
     def __call__(self, fun):
-        tools[self.name] = (fun, self.required)
+        self.fun = fun
+        tools[self.name] = self
         return fun
 
+    @property
+    def installed(self) -> bool:
+        return bool(which(self.name)) if self.external else True
+
+    @property
+    def executable(self) -> Optional[Executable]:
+        return which(self.name) if self.external else None
+
+
+#: tools we run in spack style
+tools: Dict[str, tool] = {}
+
 
 def changed_files(base="develop", untracked=True, all_files=False, root=None):
     """Get list of changed files in the Spack repository.
@@ -176,22 +189,22 @@ def setup_parser(subparser):
         "-t",
         "--tool",
         action="append",
-        help="specify which tools to run (default: %s)" % ",".join(tool_names),
+        help="specify which tools to run (default: %s)" % ", ".join(tool_names),
     )
     tool_group.add_argument(
         "-s",
         "--skip",
         metavar="TOOL",
         action="append",
-        help="specify tools to skip (choose from %s)" % ",".join(tool_names),
+        help="specify tools to skip (choose from %s)" % ", ".join(tool_names),
     )
 
     subparser.add_argument("files", nargs=argparse.REMAINDER, help="specific files to check")
 
 
-def cwd_relative(path, args):
+def cwd_relative(path, root, initial_working_dir):
     """Translate prefix-relative path to current working directory-relative."""
-    return os.path.relpath(os.path.join(args.root, path), args.initial_working_dir)
+    return os.path.relpath(os.path.join(root, path), initial_working_dir)
 
 
 def rewrite_and_print_output(
@@ -201,7 +214,10 @@ def rewrite_and_print_output(
 
     # print results relative to current working directory
     def translate(match):
-        return replacement.format(cwd_relative(match.group(1), args), *list(match.groups()[1:]))
+        return replacement.format(
+            cwd_relative(match.group(1), args.root, args.initial_working_dir),
+            *list(match.groups()[1:]),
+        )
 
     for line in output.split("\n"):
         if not line:
@@ -220,7 +236,7 @@ def print_style_header(file_list, args, tools_to_run):
     # translate modified paths to cwd_relative if needed
     paths = [filename.strip() for filename in file_list]
     if not args.root_relative:
-        paths = [cwd_relative(filename, args) for filename in paths]
+        paths = [cwd_relative(filename, args.root, args.initial_working_dir) for filename in paths]
 
     tty.msg("Modified files", *paths)
     sys.stdout.flush()
@@ -352,17 +368,127 @@ def run_black(black_cmd, file_list, args):
     return returncode
 
 
+def _module_part(root: str, expr: str):
+    parts = expr.split(".")
+    while parts:
+        f1 = os.path.join(root, "lib", "spack", *parts) + ".py"
+        f2 = os.path.join(root, "lib", "spack", *parts, "__init__.py")
+        if os.path.exists(f1) or os.path.exists(f2):
+            return ".".join(parts)
+        parts.pop()
+    return None
+
+
+def _run_import_check(
+    file_list: List[str],
+    *,
+    fix: bool,
+    root_relative: bool,
+    root=spack.paths.prefix,
+    working_dir=spack.paths.prefix,
+    out=sys.stdout,
+):
+    if sys.version_info < (3, 9):
+        print("import-check requires Python 3.9 or later")
+        return 0
+
+    is_use = re.compile(r"(?<!from )(?<!import )(?:llnl|spack)\.[a-zA-Z0-9_\.]+")
+
+    # redundant imports followed by a `# comment` are ignored, cause there can be legimitate reason
+    # to import a module: execute module scope init code, or to deal with circular imports.
+    is_abs_import = re.compile(r"^import ((?:llnl|spack)\.[a-zA-Z0-9_\.]+)$", re.MULTILINE)
+
+    exit_code = 0
+
+    for file in file_list:
+        to_add = set()
+        to_remove = []
+
+        pretty_path = file if root_relative else cwd_relative(file, root, working_dir)
+
+        try:
+            with open(file, "r") as f:
+                contents = open(file, "r").read()
+            parsed = ast.parse(contents)
+        except Exception:
+            exit_code = 1
+            print(f"{pretty_path}: could not parse", file=out)
+            continue
+
+        for m in is_abs_import.finditer(contents):
+            if contents.count(m.group(1)) == 1:
+                to_remove.append(m.group(0))
+                exit_code = 1
+                print(f"{pretty_path}: redundant import: {m.group(1)}", file=out)
+
+        # Clear all strings to avoid matching comments/strings etc.
+        for node in ast.walk(parsed):
+            if isinstance(node, ast.Constant) and isinstance(node.value, str):
+                node.value = ""
+
+        filtered_contents = ast.unparse(parsed)  # novermin
+        for m in is_use.finditer(filtered_contents):
+            module = _module_part(root, m.group(0))
+            if not module or module in to_add:
+                continue
+            if f"import {module}" not in filtered_contents:
+                to_add.add(module)
+                exit_code = 1
+                print(f"{pretty_path}: missing import: {module}", file=out)
+
+        if not fix or not to_add and not to_remove:
+            continue
+
+        with open(file, "r") as f:
+            lines = f.readlines()
+
+        if to_add:
+            # insert missing imports before the first import, delegate ordering to isort
+            for node in parsed.body:
+                if isinstance(node, (ast.Import, ast.ImportFrom)):
+                    first_line = node.lineno
+                    break
+            else:
+                print(f"{pretty_path}: could not fix", file=out)
+                continue
+            lines.insert(first_line, "\n".join(f"import {x}" for x in to_add) + "\n")
+
+        new_contents = "".join(lines)
+
+        # remove redundant imports
+        for statement in to_remove:
+            new_contents = new_contents.replace(f"{statement}\n", "")
+
+        with open(file, "w") as f:
+            f.write(new_contents)
+
+    return exit_code
+
+
+@tool("import-check", external=False)
+def run_import_check(import_check_cmd, file_list, args):
+    exit_code = _run_import_check(
+        file_list,
+        fix=args.fix,
+        root_relative=args.root_relative,
+        root=args.root,
+        working_dir=args.initial_working_dir,
+    )
+    print_tool_result("import-check", exit_code)
+    return exit_code
+
+
 def validate_toolset(arg_value):
     """Validate --tool and --skip arguments (sets of optionally comma-separated tools)."""
     tools = set(",".join(arg_value).split(","))  # allow args like 'isort,flake8'
     for tool in tools:
         if tool not in tool_names:
-            tty.die("Invaild tool: '%s'" % tool, "Choose from: %s" % ", ".join(tool_names))
+            tty.die("Invalid tool: '%s'" % tool, "Choose from: %s" % ", ".join(tool_names))
     return tools
 
 
-def missing_tools(tools_to_run):
-    return [t for t in tools_to_run if which(t) is None]
+def missing_tools(tools_to_run: List[str]) -> List[str]:
+    return [t for t in tools_to_run if not tools[t].installed]
 
 
 def _bootstrap_dev_dependencies():
@@ -417,9 +543,9 @@ def style(parser, args):
 
         print_style_header(file_list, args, tools_to_run)
         for tool_name in tools_to_run:
-            run_function, required = tools[tool_name]
+            tool = tools[tool_name]
             print_tool_header(tool_name)
-            return_code |= run_function(which(tool_name), file_list, args)
+            return_code |= tool.fun(tool.executable, file_list, args)
 
     if return_code == 0:
         tty.msg(color.colorize("@*{spack style checks were clean}"))
diff --git a/lib/spack/spack/test/cmd/style.py b/lib/spack/spack/test/cmd/style.py
index 208e31f8a2..eab4de18dd 100644
--- a/lib/spack/spack/test/cmd/style.py
+++ b/lib/spack/spack/test/cmd/style.py
@@ -4,8 +4,11 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 import filecmp
+import io
 import os
+import pathlib
 import shutil
+import sys
 
 import pytest
 
@@ -15,7 +18,7 @@ import spack.cmd.style
 import spack.main
 import spack.paths
 import spack.repo
-from spack.cmd.style import changed_files
+from spack.cmd.style import _run_import_check, changed_files
 from spack.util.executable import which
 
 #: directory with sample style files
@@ -292,5 +295,97 @@ def test_style_with_black(flake8_package_with_errors):
 
 
 def test_skip_tools():
-    output = style("--skip", "isort,mypy,black,flake8")
+    output = style("--skip", "import-check,isort,mypy,black,flake8")
     assert "Nothing to run" in output
+
+
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="requires Python 3.9+")
+def test_run_import_check(tmp_path: pathlib.Path):
+    file = tmp_path / "issues.py"
+    contents = '''
+import spack.cmd
+import spack.config  # do not drop this import because of this comment
+
+# this comment about spack.error should not be removed
+class Example(spack.build_systems.autotools.AutotoolsPackage):
+    """this is a docstring referencing unused spack.error.SpackError, which is fine"""
+    pass
+
+def foo(config: "spack.error.SpackError"):
+    # the type hint is quoted, so it should not be removed
+    spack.util.executable.Executable("example")
+'''
+    file.write_text(contents)
+    root = str(tmp_path)
+    output_buf = io.StringIO()
+    exit_code = _run_import_check(
+        [str(file)],
+        fix=False,
+        out=output_buf,
+        root_relative=False,
+        root=spack.paths.prefix,
+        working_dir=root,
+    )
+    output = output_buf.getvalue()
+
+    assert "issues.py: redundant import: spack.cmd" in output
+    assert "issues.py: redundant import: spack.config" not in output  # comment prevents removal
+    assert "issues.py: missing import: spack.build_systems.autotools" in output
+    assert "issues.py: missing import: spack.util.executable" in output
+    assert "issues.py: missing import: spack.error" not in output  # not directly used
+    assert exit_code == 1
+    assert file.read_text() == contents  # fix=False should not change the file
+
+    # run it with --fix, should have the same output.
+    output_buf = io.StringIO()
+    exit_code = _run_import_check(
+        [str(file)],
+        fix=True,
+        out=output_buf,
+        root_relative=False,
+        root=spack.paths.prefix,
+        working_dir=root,
+    )
+    output = output_buf.getvalue()
+    assert exit_code == 1
+    assert "issues.py: redundant import: spack.cmd" in output
+    assert "issues.py: missing import: spack.build_systems.autotools" in output
+    assert "issues.py: missing import: spack.util.executable" in output
+
+    # after fix a second fix is idempotent
+    output_buf = io.StringIO()
+    exit_code = _run_import_check(
+        [str(file)],
+        fix=True,
+        out=output_buf,
+        root_relative=False,
+        root=spack.paths.prefix,
+        working_dir=root,
+    )
+    output = output_buf.getvalue()
+    assert exit_code == 0
+    assert not output
+
+    # check that the file was fixed
+    new_contents = file.read_text()
+    assert "import spack.cmd" not in new_contents
+    assert "import spack.build_systems.autotools" in new_contents
+    assert "import spack.util.executable" in new_contents
+
+
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="requires Python 3.9+")
+def test_run_import_check_syntax_error_and_missing(tmp_path: pathlib.Path):
+    (tmp_path / "syntax-error.py").write_text("""this 'is n(ot python code""")
+    output_buf = io.StringIO()
+    exit_code = _run_import_check(
+        [str(tmp_path / "syntax-error.py"), str(tmp_path / "missing.py")],
+        fix=False,
+        out=output_buf,
+        root_relative=True,
+        root=str(tmp_path),
+        working_dir=str(tmp_path / "does-not-matter"),
+    )
+    output = output_buf.getvalue()
+    assert "syntax-error.py: could not parse" in output
+    assert "missing.py: could not parse" in output
+    assert exit_code == 1
diff --git a/share/spack/spack-completion.fish b/share/spack/spack-completion.fish
index 385361389a..6c8f2521b4 100644
--- a/share/spack/spack-completion.fish
+++ b/share/spack/spack-completion.fish
@@ -2908,9 +2908,9 @@ complete -c spack -n '__fish_spack_using_command style' -s f -l fix -d 'format a
 complete -c spack -n '__fish_spack_using_command style' -l root -r -f -a root
 complete -c spack -n '__fish_spack_using_command style' -l root -r -d 'style check a different spack instance'
 complete -c spack -n '__fish_spack_using_command style' -s t -l tool -r -f -a tool
-complete -c spack -n '__fish_spack_using_command style' -s t -l tool -r -d 'specify which tools to run (default: isort,black,flake8,mypy)'
+complete -c spack -n '__fish_spack_using_command style' -s t -l tool -r -d 'specify which tools to run (default: import-check, isort, black, flake8, mypy)'
 complete -c spack -n '__fish_spack_using_command style' -s s -l skip -r -f -a skip
-complete -c spack -n '__fish_spack_using_command style' -s s -l skip -r -d 'specify tools to skip (choose from isort,black,flake8,mypy)'
+complete -c spack -n '__fish_spack_using_command style' -s s -l skip -r -d 'specify tools to skip (choose from import-check, isort, black, flake8, mypy)'
 
 # spack tags
 set -g __fish_spack_optspecs_spack_tags h/help i/installed a/all