spack repo migrate: support v1 -> v2 migration for package api (#50507)

lib/spack/spack/cmd/repo.py

@@ -2,10 +2,10 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-import ast
 import os
+import shlex
 import sys
-from typing import Any, Dict, List, Optional, Set
+from typing import Any, List, Optional
 
 import llnl.util.tty as tty
 
@@ -72,6 +72,9 @@ def setup_parser(subparser):
     migrate_parser.add_argument(
         "namespace_or_path", help="path to a Spack package repository directory"
     )
+    migrate_parser.add_argument(
+        "--fix", action="store_true", help="automatically fix the imports in the package files"
+    )
 
 
 def repo_create(args):
@@ -171,204 +174,62 @@ def _get_repo(name_or_path: str) -> Optional[spack.repo.Repo]:
 
     for repo in spack.config.get("repos"):
         try:
-            r = spack.repo.from_path(spack.util.path.canonicalize_path(repo))
+            r = spack.repo.from_path(repo)
         except spack.repo.RepoError:
             continue
-        if r.namespace == name_or_path or os.path.samefile(r.root, name_or_path):
+        if r.namespace == name_or_path:
             return r
-
     return None
 
 
-def repo_migrate(args: Any) -> None:
+def repo_migrate(args: Any) -> int:
     """migrate a package repository to the latest Package API"""
+    from spack.repo_migrate import migrate_v1_to_v2, migrate_v2_imports
+
     repo = _get_repo(args.namespace_or_path)
 
     if repo is None:
         tty.die(f"No such repository: {args.namespace_or_path}")
 
-    if repo.package_api < (2, 0):
+    if (1, 0) <= repo.package_api < (2, 0):
-        tty.die("Migration from Spack repo API < 2.0 is not supported yet")
+        success, repo_v2 = migrate_v1_to_v2(repo, fix=args.fix)
+        exit_code = 0 if success else 1
+    elif (2, 0) <= repo.package_api < (3, 0):
+        repo_v2 = None
+        exit_code = 0 if migrate_v2_imports(repo.packages_path, repo.root, fix=args.fix) else 1
+    else:
+        repo_v2 = None
+        exit_code = 0
 
-    symbol_to_module = {
+    if exit_code == 0 and isinstance(repo_v2, spack.repo.Repo):
-        "AspellDictPackage": "spack.build_systems.aspell_dict",
+        tty.info(
-        "AutotoolsPackage": "spack.build_systems.autotools",
+            f"Repository '{repo_v2.namespace}' was successfully migrated from "
-        "BundlePackage": "spack.build_systems.bundle",
+            f"package API {repo.package_api_str} to {repo_v2.package_api_str}."
-        "CachedCMakePackage": "spack.build_systems.cached_cmake",
-        "cmake_cache_filepath": "spack.build_systems.cached_cmake",
-        "cmake_cache_option": "spack.build_systems.cached_cmake",
-        "cmake_cache_path": "spack.build_systems.cached_cmake",
-        "cmake_cache_string": "spack.build_systems.cached_cmake",
-        "CargoPackage": "spack.build_systems.cargo",
-        "CMakePackage": "spack.build_systems.cmake",
-        "generator": "spack.build_systems.cmake",
-        "CompilerPackage": "spack.build_systems.compiler",
-        "CudaPackage": "spack.build_systems.cuda",
-        "Package": "spack.build_systems.generic",
-        "GNUMirrorPackage": "spack.build_systems.gnu",
-        "GoPackage": "spack.build_systems.go",
-        "IntelPackage": "spack.build_systems.intel",
-        "LuaPackage": "spack.build_systems.lua",
-        "MakefilePackage": "spack.build_systems.makefile",
-        "MavenPackage": "spack.build_systems.maven",
-        "MesonPackage": "spack.build_systems.meson",
-        "MSBuildPackage": "spack.build_systems.msbuild",
-        "NMakePackage": "spack.build_systems.nmake",
-        "OctavePackage": "spack.build_systems.octave",
-        "INTEL_MATH_LIBRARIES": "spack.build_systems.oneapi",
-        "IntelOneApiLibraryPackage": "spack.build_systems.oneapi",
-        "IntelOneApiLibraryPackageWithSdk": "spack.build_systems.oneapi",
-        "IntelOneApiPackage": "spack.build_systems.oneapi",
-        "IntelOneApiStaticLibraryList": "spack.build_systems.oneapi",
-        "PerlPackage": "spack.build_systems.perl",
-        "PythonExtension": "spack.build_systems.python",
-        "PythonPackage": "spack.build_systems.python",
-        "QMakePackage": "spack.build_systems.qmake",
-        "RPackage": "spack.build_systems.r",
-        "RacketPackage": "spack.build_systems.racket",
-        "ROCmPackage": "spack.build_systems.rocm",
-        "RubyPackage": "spack.build_systems.ruby",
-        "SConsPackage": "spack.build_systems.scons",
-        "SIPPackage": "spack.build_systems.sip",
-        "SourceforgePackage": "spack.build_systems.sourceforge",
-        "SourcewarePackage": "spack.build_systems.sourceware",
-        "WafPackage": "spack.build_systems.waf",
-        "XorgPackage": "spack.build_systems.xorg",
-    }
-
-    for f in os.scandir(repo.packages_path):
-        pkg_path = os.path.join(f.path, "package.py")
-        try:
-            if f.name in ("__init__.py", "__pycache__") or not f.is_dir():
-                continue
-            with open(pkg_path, "rb") as file:
-                tree = ast.parse(file.read())
-        except (OSError, SyntaxError) as e:
-            print(f"Skipping {pkg_path}: {e}", file=sys.stderr)
-            continue
-
-        #: Symbols that are referenced in the package and may need to be imported.
-        referenced_symbols: Set[str] = set()
-
-        #: Set of symbols of interest that are already defined through imports, assignments, or
-        #: function definitions.
-        defined_symbols: Set[str] = set()
-
-        best_line: Optional[int] = None
-
-        seen_import = False
-
-        for node in ast.walk(tree):
-            # Get the last import statement from the first block of top-level imports
-            if isinstance(node, ast.Module):
-                for child in ast.iter_child_nodes(node):
-                    # if we never encounter an import statement, the best line to add is right
-                    # before the first node under the module
-                    if best_line is None and isinstance(child, ast.stmt):
-                        best_line = child.lineno
-
-                    # prefer adding right before `from spack.package import ...`
-                    if isinstance(child, ast.ImportFrom) and child.module == "spack.package":
-                        seen_import = True
-                        best_line = child.lineno  # add it right before spack.package
-                        break
-
-                    # otherwise put it right after the last import statement
-                    is_import = isinstance(child, (ast.Import, ast.ImportFrom))
-
-                    if is_import:
-                        if isinstance(child, (ast.stmt, ast.expr)):
-                            best_line = (child.end_lineno or child.lineno) + 1
-
-                    if not seen_import and is_import:
-                        seen_import = True
-                    elif seen_import and not is_import:
-                        break
-
-            # Function definitions or assignments to variables whose name is a symbol of interest
-            # are considered as redefinitions, so we skip them.
-            elif isinstance(node, ast.FunctionDef):
-                if node.name in symbol_to_module:
-                    print(
-                        f"{pkg_path}:{node.lineno}: redefinition of `{node.name}` skipped",
-                        file=sys.stderr,
         )
-                    defined_symbols.add(node.name)
+        tty.warn(
-            elif isinstance(node, ast.Assign):
+            "Remove the old repository from Spack's configuration and add the new one using:\n"
-                for target in node.targets:
+            f"    spack repo remove {shlex.quote(repo.root)}\n"
-                    if isinstance(target, ast.Name) and target.id in symbol_to_module:
+            f"    spack repo add {shlex.quote(repo_v2.root)}"
-                        print(
-                            f"{pkg_path}:{target.lineno}: redefinition of `{target.id}` skipped",
-                            file=sys.stderr,
-                        )
-                        defined_symbols.add(target.id)
-
-            # Register symbols that are not imported.
-            elif isinstance(node, ast.Name) and node.id in symbol_to_module:
-                referenced_symbols.add(node.id)
-
-            # Register imported symbols to make this operation idempotent
-            elif isinstance(node, ast.ImportFrom):
-                for alias in node.names:
-                    if alias.name in symbol_to_module:
-                        defined_symbols.add(alias.name)
-                        if node.module == "spack.package":
-                            print(
-                                f"{pkg_path}:{node.lineno}: `{alias.name}` is imported from "
-                                "`spack.package`, which no longer provides this symbol",
-                                file=sys.stderr,
         )
 
-                    if alias.asname and alias.asname in symbol_to_module:
+    elif exit_code == 0:
-                        defined_symbols.add(alias.asname)
+        tty.info(f"Repository '{repo.namespace}' was successfully migrated")
 
-        # Remove imported symbols from the referenced symbols
+    elif not args.fix and exit_code == 1:
-        referenced_symbols.difference_update(defined_symbols)
+        tty.error(
+            f"No changes were made to the repository {repo.root} with namespace "
+            f"'{repo.namespace}'. Run with --fix to apply the above changes."
+        )
 
-        if not referenced_symbols:
+    return exit_code
-            continue
-
-        if best_line is None:
-            print(f"{pkg_path}: failed to update imports", file=sys.stderr)
-            continue
-
-        # Add the missing imports right after the last import statement
-        with open(pkg_path, "r", encoding="utf-8") as file:
-            lines = file.readlines()
-
-        # Group missing symbols by their module
-        missing_imports_by_module: Dict[str, list] = {}
-        for symbol in referenced_symbols:
-            module = symbol_to_module[symbol]
-            if module not in missing_imports_by_module:
-                missing_imports_by_module[module] = []
-            missing_imports_by_module[module].append(symbol)
-
-        new_lines = [
-            f"from {module} import {', '.join(sorted(symbols))}\n"
-            for module, symbols in sorted(missing_imports_by_module.items())
-        ]
-
-        if not seen_import:
-            new_lines.extend(("\n", "\n"))
-
-        lines[best_line - 1 : best_line - 1] = new_lines
-
-        tmp_file = pkg_path + ".tmp"
-
-        with open(tmp_file, "w", encoding="utf-8") as file:
-            file.writelines(lines)
-
-        os.replace(tmp_file, pkg_path)
 
 
 def repo(parser, args):
-    action = {
+    return {
         "create": repo_create,
         "list": repo_list,
         "add": repo_add,
         "remove": repo_remove,
         "rm": repo_remove,
         "migrate": repo_migrate,
-    }
+    }[args.repo_command](args)
-    action[args.repo_command](args)

lib/spack/spack/repo_migrate.py

@@ -0,0 +1,421 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+import ast
+import os
+import re
+import shutil
+import sys
+from typing import IO, Dict, List, Optional, Set, Tuple
+
+import spack.repo
+import spack.util.naming
+import spack.util.spack_yaml
+
+
+def _same_contents(f: str, g: str) -> bool:
+    """Return True if the files have the same contents."""
+    try:
+        with open(f, "rb") as f1, open(g, "rb") as f2:
+            while True:
+                b1 = f1.read(4096)
+                b2 = f2.read(4096)
+                if b1 != b2:
+                    return False
+                if not b1 and not b2:
+                    break
+            return True
+    except OSError:
+        return False
+
+
+def migrate_v1_to_v2(
+    repo: spack.repo.Repo, fix: bool, out: IO[str] = sys.stdout, err: IO[str] = sys.stderr
+) -> Tuple[bool, Optional[spack.repo.Repo]]:
+    """To upgrade a repo from Package API v1 to v2 we need to:
+    1. ensure ``spack_repo/<namespace>`` parent dirs to the ``repo.yaml`` file.
+    2. rename <pkg dir>/package.py to <pkg module>/package.py.
+    3. bump the version in ``repo.yaml``.
+    """
+    if not (1, 0) <= repo.package_api < (2, 0):
+        raise RuntimeError(f"Cannot upgrade from {repo.package_api_str} to v2.0")
+
+    with open(os.path.join(repo.root, "repo.yaml"), encoding="utf-8") as f:
+        updated_config = spack.util.spack_yaml.load(f)
+        updated_config["repo"]["api"] = "v2.0"
+
+    namespace = repo.namespace.split(".")
+
+    if not all(
+        spack.util.naming.valid_module_name(part, package_api=(2, 0)) for part in namespace
+    ):
+        print(
+            f"Cannot upgrade from v1 to v2, because the namespace '{repo.namespace}' is not a "
+            "valid Python module",
+            file=err,
+        )
+        return False, None
+
+    try:
+        subdirectory = spack.repo._validate_and_normalize_subdir(
+            repo.subdirectory, repo.root, package_api=(2, 0)
+        )
+    except spack.repo.BadRepoError:
+        print(
+            f"Cannot upgrade from v1 to v2, because the subdirectory '{repo.subdirectory}' is not "
+            "a valid Python module",
+            file=err,
+        )
+        return False, None
+
+    new_root = os.path.join(repo.root, "spack_repo", *namespace)
+
+    ino_to_relpath: Dict[int, str] = {}
+    symlink_to_ino: Dict[str, int] = {}
+
+    prefix_len = len(repo.root) + len(os.sep)
+
+    rename: Dict[str, str] = {}
+    dirs_to_create: List[str] = []
+    files_to_copy: List[str] = []
+
+    errors = False
+
+    stack: List[Tuple[str, int]] = [(repo.root, 0)]
+    while stack:
+        path, depth = stack.pop()
+
+        try:
+            entries = os.scandir(path)
+        except OSError:
+            continue
+
+        for entry in entries:
+            rel_path = entry.path[prefix_len:]
+
+            if depth == 0 and entry.name in ("spack_repo", "repo.yaml"):
+                continue
+
+            ino_to_relpath[entry.inode()] = entry.path[prefix_len:]
+
+            if entry.is_symlink():
+                symlink_to_ino[rel_path] = entry.stat(follow_symlinks=True).st_ino
+                continue
+
+            elif entry.is_dir(follow_symlinks=False):
+                if entry.name == "__pycache__":
+                    continue
+
+                # check if this is a package
+                if (
+                    depth == 1
+                    and rel_path.startswith(f"{subdirectory}{os.sep}")
+                    and os.path.exists(os.path.join(entry.path, "package.py"))
+                ):
+                    if "_" in entry.name:
+                        print(
+                            f"Invalid package name '{entry.name}': underscores are not allowed in "
+                            "package names, rename the package with hyphens as separators",
+                            file=err,
+                        )
+                        errors = True
+                        continue
+                    pkg_dir = spack.util.naming.pkg_name_to_pkg_dir(entry.name, package_api=(2, 0))
+                    if pkg_dir != entry.name:
+                        rename[f"{subdirectory}{os.sep}{entry.name}"] = (
+                            f"{subdirectory}{os.sep}{pkg_dir}"
+                        )
+
+                dirs_to_create.append(rel_path)
+
+                stack.append((entry.path, depth + 1))
+                continue
+
+            files_to_copy.append(rel_path)
+
+    if errors:
+        return False, None
+
+    rename_regex = re.compile("^(" + "|".join(re.escape(k) for k in rename.keys()) + ")")
+
+    if fix:
+        os.makedirs(new_root, exist_ok=True)
+
+    def _relocate(rel_path: str) -> Tuple[str, str]:
+        return os.path.join(repo.root, rel_path), os.path.join(
+            new_root, rename_regex.sub(lambda m: rename[m.group(0)], rel_path)
+        )
+
+    if not fix:
+        print("The following directories, files and symlinks will be created:\n", file=out)
+
+    for rel_path in dirs_to_create:
+        _, new_path = _relocate(rel_path)
+        if fix:
+            try:
+                os.mkdir(new_path)
+            except FileExistsError:  # not an error if the directory already exists
+                continue
+        else:
+            print(f"create directory {new_path}", file=out)
+
+    for rel_path in files_to_copy:
+        old_path, new_path = _relocate(rel_path)
+        if os.path.lexists(new_path):
+            # if we already copied this file, don't error.
+            if not _same_contents(old_path, new_path):
+                print(
+                    f"Cannot upgrade from v1 to v2, because the file '{new_path}' already exists",
+                    file=err,
+                )
+                return False, None
+            continue
+        if fix:
+            shutil.copy2(old_path, new_path)
+        else:
+            print(f"copy {old_path} -> {new_path}", file=out)
+
+    for rel_path, ino in symlink_to_ino.items():
+        old_path, new_path = _relocate(rel_path)
+        if ino in ino_to_relpath:
+            # link by path relative to the new root
+            _, new_target = _relocate(ino_to_relpath[ino])
+            tgt = os.path.relpath(new_target, new_path)
+        else:
+            tgt = os.path.realpath(old_path)
+
+        # no-op if the same, error if different
+        if os.path.lexists(new_path):
+            if not os.path.islink(new_path) or os.readlink(new_path) != tgt:
+                print(
+                    f"Cannot upgrade from v1 to v2, because the file '{new_path}' already exists",
+                    file=err,
+                )
+                return False, None
+            continue
+
+        if fix:
+            os.symlink(tgt, new_path)
+        else:
+            print(f"create symlink {new_path} -> {tgt}", file=out)
+
+    if fix:
+        with open(os.path.join(new_root, "repo.yaml"), "w", encoding="utf-8") as f:
+            spack.util.spack_yaml.dump(updated_config, f)
+        updated_repo = spack.repo.from_path(new_root)
+    else:
+        print(file=out)
+        updated_repo = repo  # compute the import diff on the v1 repo since v2 doesn't exist yet
+
+    result = migrate_v2_imports(
+        updated_repo.packages_path, updated_repo.root, fix=fix, out=out, err=err
+    )
+
+    return result, (updated_repo if fix else None)
+
+
+def migrate_v2_imports(
+    packages_dir: str, root: str, fix: bool, out: IO[str] = sys.stdout, err: IO[str] = sys.stderr
+) -> bool:
+    """In Package API v2.0, packages need to explicitly import package classes and a few other
+    symbols from the build_systems module. This function automatically adds the missing imports
+    to each package.py file in the repository."""
+
+    symbol_to_module = {
+        "AspellDictPackage": "spack_repo.builtin.build_systems.aspell_dict",
+        "AutotoolsPackage": "spack_repo.builtin.build_systems.autotools",
+        "BundlePackage": "spack_repo.builtin.build_systems.bundle",
+        "CachedCMakePackage": "spack_repo.builtin.build_systems.cached_cmake",
+        "cmake_cache_filepath": "spack_repo.builtin.build_systems.cached_cmake",
+        "cmake_cache_option": "spack_repo.builtin.build_systems.cached_cmake",
+        "cmake_cache_path": "spack_repo.builtin.build_systems.cached_cmake",
+        "cmake_cache_string": "spack_repo.builtin.build_systems.cached_cmake",
+        "CargoPackage": "spack_repo.builtin.build_systems.cargo",
+        "CMakePackage": "spack_repo.builtin.build_systems.cmake",
+        "generator": "spack_repo.builtin.build_systems.cmake",
+        "CompilerPackage": "spack_repo.builtin.build_systems.compiler",
+        "CudaPackage": "spack_repo.builtin.build_systems.cuda",
+        "Package": "spack_repo.builtin.build_systems.generic",
+        "GNUMirrorPackage": "spack_repo.builtin.build_systems.gnu",
+        "GoPackage": "spack_repo.builtin.build_systems.go",
+        "IntelPackage": "spack_repo.builtin.build_systems.intel",
+        "LuaPackage": "spack_repo.builtin.build_systems.lua",
+        "MakefilePackage": "spack_repo.builtin.build_systems.makefile",
+        "MavenPackage": "spack_repo.builtin.build_systems.maven",
+        "MesonPackage": "spack_repo.builtin.build_systems.meson",
+        "MSBuildPackage": "spack_repo.builtin.build_systems.msbuild",
+        "NMakePackage": "spack_repo.builtin.build_systems.nmake",
+        "OctavePackage": "spack_repo.builtin.build_systems.octave",
+        "INTEL_MATH_LIBRARIES": "spack_repo.builtin.build_systems.oneapi",
+        "IntelOneApiLibraryPackage": "spack_repo.builtin.build_systems.oneapi",
+        "IntelOneApiLibraryPackageWithSdk": "spack_repo.builtin.build_systems.oneapi",
+        "IntelOneApiPackage": "spack_repo.builtin.build_systems.oneapi",
+        "IntelOneApiStaticLibraryList": "spack_repo.builtin.build_systems.oneapi",
+        "PerlPackage": "spack_repo.builtin.build_systems.perl",
+        "PythonExtension": "spack_repo.builtin.build_systems.python",
+        "PythonPackage": "spack_repo.builtin.build_systems.python",
+        "QMakePackage": "spack_repo.builtin.build_systems.qmake",
+        "RPackage": "spack_repo.builtin.build_systems.r",
+        "RacketPackage": "spack_repo.builtin.build_systems.racket",
+        "ROCmPackage": "spack_repo.builtin.build_systems.rocm",
+        "RubyPackage": "spack_repo.builtin.build_systems.ruby",
+        "SConsPackage": "spack_repo.builtin.build_systems.scons",
+        "SIPPackage": "spack_repo.builtin.build_systems.sip",
+        "SourceforgePackage": "spack_repo.builtin.build_systems.sourceforge",
+        "SourcewarePackage": "spack_repo.builtin.build_systems.sourceware",
+        "WafPackage": "spack_repo.builtin.build_systems.waf",
+        "XorgPackage": "spack_repo.builtin.build_systems.xorg",
+    }
+
+    success = True
+
+    for f in os.scandir(packages_dir):
+        pkg_path = os.path.join(f.path, "package.py")
+        if (
+            f.name in ("__init__.py", "__pycache__")
+            or not f.is_dir(follow_symlinks=False)
+            or os.path.islink(pkg_path)
+        ):
+            print(f"Skipping {f.path}", file=err)
+            continue
+        try:
+            with open(pkg_path, "rb") as file:
+                tree = ast.parse(file.read())
+        except (OSError, SyntaxError) as e:
+            print(f"Skipping {pkg_path}: {e}", file=err)
+            continue
+
+        #: Symbols that are referenced in the package and may need to be imported.
+        referenced_symbols: Set[str] = set()
+
+        #: Set of symbols of interest that are already defined through imports, assignments, or
+        #: function definitions.
+        defined_symbols: Set[str] = set()
+
+        best_line: Optional[int] = None
+
+        seen_import = False
+
+        for node in ast.walk(tree):
+            # Get the last import statement from the first block of top-level imports
+            if isinstance(node, ast.Module):
+                for child in ast.iter_child_nodes(node):
+                    # if we never encounter an import statement, the best line to add is right
+                    # before the first node under the module
+                    if best_line is None and isinstance(child, ast.stmt):
+                        best_line = child.lineno
+
+                    # prefer adding right before `from spack.package import ...`
+                    if isinstance(child, ast.ImportFrom) and child.module == "spack.package":
+                        seen_import = True
+                        best_line = child.lineno  # add it right before spack.package
+                        break
+
+                    # otherwise put it right after the last import statement
+                    is_import = isinstance(child, (ast.Import, ast.ImportFrom))
+
+                    if is_import:
+                        if isinstance(child, (ast.stmt, ast.expr)):
+                            best_line = (child.end_lineno or child.lineno) + 1
+
+                    if not seen_import and is_import:
+                        seen_import = True
+                    elif seen_import and not is_import:
+                        break
+
+            # Function definitions or assignments to variables whose name is a symbol of interest
+            # are considered as redefinitions, so we skip them.
+            elif isinstance(node, ast.FunctionDef):
+                if node.name in symbol_to_module:
+                    print(
+                        f"{pkg_path}:{node.lineno}: redefinition of `{node.name}` skipped",
+                        file=err,
+                    )
+                    defined_symbols.add(node.name)
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and target.id in symbol_to_module:
+                        print(
+                            f"{pkg_path}:{target.lineno}: redefinition of `{target.id}` skipped",
+                            file=err,
+                        )
+                        defined_symbols.add(target.id)
+
+            # Register symbols that are not imported.
+            elif isinstance(node, ast.Name) and node.id in symbol_to_module:
+                referenced_symbols.add(node.id)
+
+            # Register imported symbols to make this operation idempotent
+            elif isinstance(node, ast.ImportFrom):
+                for alias in node.names:
+                    if alias.name in symbol_to_module:
+                        defined_symbols.add(alias.name)
+                        if node.module == "spack.package":
+                            success = False
+                            print(
+                                f"{pkg_path}:{node.lineno}: `{alias.name}` is imported from "
+                                "`spack.package`, which no longer provides this symbol",
+                                file=err,
+                            )
+
+                    if alias.asname and alias.asname in symbol_to_module:
+                        defined_symbols.add(alias.asname)
+
+        # Remove imported symbols from the referenced symbols
+        referenced_symbols.difference_update(defined_symbols)
+
+        if not referenced_symbols:
+            continue
+
+        if best_line is None:
+            print(f"{pkg_path}: failed to update imports", file=err)
+            success = False
+            continue
+
+        # Add the missing imports right after the last import statement
+        with open(pkg_path, "r", encoding="utf-8", newline="") as file:
+            lines = file.readlines()
+
+        # Group missing symbols by their module
+        missing_imports_by_module: Dict[str, list] = {}
+        for symbol in referenced_symbols:
+            module = symbol_to_module[symbol]
+            if module not in missing_imports_by_module:
+                missing_imports_by_module[module] = []
+            missing_imports_by_module[module].append(symbol)
+
+        new_lines = [
+            f"from {module} import {', '.join(sorted(symbols))}\n"
+            for module, symbols in sorted(missing_imports_by_module.items())
+        ]
+
+        if not seen_import:
+            new_lines.extend(("\n", "\n"))
+
+        if not fix:  # only print the diff
+            success = False  # packages need to be fixed, but we didn't do it
+            diff_start, diff_end = max(1, best_line - 3), min(best_line + 2, len(lines))
+            num_changed = diff_end - diff_start + 1
+            num_added = num_changed + len(new_lines)
+            rel_pkg_path = os.path.relpath(pkg_path, start=root)
+            out.write(f"--- a/{rel_pkg_path}\n+++ b/{rel_pkg_path}\n")
+            out.write(f"@@ -{diff_start},{num_changed} +{diff_start},{num_added} @@\n")
+            for line in lines[diff_start - 1 : best_line - 1]:
+                out.write(f" {line}")
+            for line in new_lines:
+                out.write(f"+{line}")
+            for line in lines[best_line - 1 : diff_end]:
+                out.write(f" {line}")
+            continue
+
+        lines[best_line - 1 : best_line - 1] = new_lines
+
+        tmp_file = pkg_path + ".tmp"
+
+        with open(tmp_file, "w", encoding="utf-8", newline="") as file:
+            file.writelines(lines)
+
+        os.replace(tmp_file, pkg_path)
+
+    return success

lib/spack/spack/test/cmd/repo.py

@@ -1,16 +1,21 @@
 # Copyright Spack Project Developers. See COPYRIGHT file for details.
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
+import io
 import os
 import pathlib
 
 import pytest
 
+from llnl.util.filesystem import working_dir
+
 import spack.config
 import spack.environment as ev
 import spack.main
 import spack.repo
+import spack.repo_migrate
 from spack.main import SpackCommand
+from spack.util.executable import Executable
 
 repo = spack.main.SpackCommand("repo")
 env = SpackCommand("env")
@@ -71,63 +76,98 @@ def test_env_repo_path_vars_substitution(
             assert current_dir in repos_specs
 
 
+OLD_7ZIP = b"""\
+# some comment
+
+from spack.package import *
+
+class _7zip(Package):
+    pass
+"""
+
+NEW_7ZIP = b"""\
+# some comment
+
+from spack_repo.builtin.build_systems.generic import Package
+from spack.package import *
+
+class _7zip(Package):
+    pass
+"""
+
+OLD_NUMPY = b"""\
+# some comment
+
+from spack.package import *
+
+class PyNumpy(CMakePackage):
+    generator("ninja")
+"""
+
+NEW_NUMPY = b"""\
+# some comment
+
+from spack_repo.builtin.build_systems.cmake import CMakePackage, generator
+from spack.package import *
+
+class PyNumpy(CMakePackage):
+    generator("ninja")
+"""
+
+
 def test_repo_migrate(tmp_path: pathlib.Path, config):
-    path, _ = spack.repo.create_repo(str(tmp_path), "mockrepo", package_api=(2, 0))
+    old_root, _ = spack.repo.create_repo(str(tmp_path), "org.repo", package_api=(1, 0))
-    pkgs_path = spack.repo.from_path(path).packages_path
+    pkgs_path = pathlib.Path(spack.repo.from_path(old_root).packages_path)
+    new_root = pathlib.Path(old_root) / "spack_repo" / "org" / "repo"
 
-    pkg1 = pathlib.Path(os.path.join(pkgs_path, "foo", "package.py"))
+    pkg_7zip_old = pkgs_path / "7zip" / "package.py"
-    pkg2 = pathlib.Path(os.path.join(pkgs_path, "bar", "package.py"))
+    pkg_numpy_old = pkgs_path / "py-numpy" / "package.py"
+    pkg_py_7zip_new = new_root / "packages" / "_7zip" / "package.py"
+    pkg_py_numpy_new = new_root / "packages" / "py_numpy" / "package.py"
 
-    pkg1.parent.mkdir(parents=True)
+    pkg_7zip_old.parent.mkdir(parents=True)
-    pkg2.parent.mkdir(parents=True)
+    pkg_numpy_old.parent.mkdir(parents=True)
 
-    pkg1.write_text(
+    pkg_7zip_old.write_bytes(OLD_7ZIP)
-        """\
+    pkg_numpy_old.write_bytes(OLD_NUMPY)
-# some comment
 
-from spack.package import *
+    repo("migrate", "--fix", old_root)
 
-class Foo(Package):
+    # old files are not touched since they are moved
-    pass
+    assert pkg_7zip_old.read_bytes() == OLD_7ZIP
-""",
+    assert pkg_numpy_old.read_bytes() == OLD_NUMPY
-        encoding="utf-8",
-    )
-    pkg2.write_text(
-        """\
-# some comment
 
-from spack.package import *
+    # new files are created and have updated contents
+    assert pkg_py_7zip_new.read_bytes() == NEW_7ZIP
+    assert pkg_py_numpy_new.read_bytes() == NEW_NUMPY
 
-class Bar(CMakePackage):
+
-    generator("ninja")
+def test_migrate_diff(git: Executable, tmp_path: pathlib.Path):
-""",
+    root, _ = spack.repo.create_repo(str(tmp_path), "foo", package_api=(2, 0))
-        encoding="utf-8",
+    r = pathlib.Path(root)
+    pkg_7zip = r / "packages" / "_7zip" / "package.py"
+    pkg_py_numpy_new = r / "packages" / "py_numpy" / "package.py"
+    pkg_broken = r / "packages" / "broken" / "package.py"
+
+    pkg_7zip.parent.mkdir(parents=True)
+    pkg_py_numpy_new.parent.mkdir(parents=True)
+    pkg_broken.parent.mkdir(parents=True)
+    pkg_7zip.write_bytes(OLD_7ZIP)
+    pkg_py_numpy_new.write_bytes(OLD_NUMPY)
+    pkg_broken.write_bytes(b"syntax(error")
+
+    stderr = io.StringIO()
+
+    with open(tmp_path / "imports.patch", "w", encoding="utf-8") as stdout:
+        spack.repo_migrate.migrate_v2_imports(
+            str(r / "packages"), str(r), fix=False, out=stdout, err=stderr
         )
 
-    repo("migrate", path)
+    assert f"Skipping {pkg_broken}" in stderr.getvalue()
 
-    assert (
+    # apply the patch and verify the changes
-        pkg1.read_text(encoding="utf-8")
+    with working_dir(str(r)):
-        == """\
+        git("apply", str(tmp_path / "imports.patch"))
-# some comment
 
-from spack.build_systems.generic import Package
+    assert pkg_7zip.read_bytes() == NEW_7ZIP
-from spack.package import *
+    assert pkg_py_numpy_new.read_bytes() == NEW_NUMPY
-
-class Foo(Package):
-    pass
-"""
-    )
-
-    assert (
-        pkg2.read_text(encoding="utf-8")
-        == """\
-# some comment
-
-from spack.build_systems.cmake import CMakePackage, generator
-from spack.package import *
-
-class Bar(CMakePackage):
-    generator("ninja")
-"""
-    )

share/spack/spack-completion.bash

@@ -1831,7 +1831,7 @@ _spack_repo_rm() {
 _spack_repo_migrate() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help"
+        SPACK_COMPREPLY="-h --help --fix"
     else
         _repos
     fi

share/spack/spack-completion.fish

@@ -2793,10 +2793,12 @@ complete -c spack -n '__fish_spack_using_command repo rm' -l scope -r -f -a '_bu
 complete -c spack -n '__fish_spack_using_command repo rm' -l scope -r -d 'configuration scope to modify'
 
 # spack repo migrate
-set -g __fish_spack_optspecs_spack_repo_migrate h/help
+set -g __fish_spack_optspecs_spack_repo_migrate h/help fix
 complete -c spack -n '__fish_spack_using_command_pos 0 repo migrate' $__fish_spack_force_files -a '(__fish_spack_repos)'
 complete -c spack -n '__fish_spack_using_command repo migrate' -s h -l help -f -a help
 complete -c spack -n '__fish_spack_using_command repo migrate' -s h -l help -d 'show this help message and exit'
+complete -c spack -n '__fish_spack_using_command repo migrate' -l fix -f -a fix
+complete -c spack -n '__fish_spack_using_command repo migrate' -l fix -d 'automatically fix the imports in the package files'
 
 # spack resource
 set -g __fish_spack_optspecs_spack_resource h/help