Use spack.traverse.traverse_nodes where useful (#33677)

1 files changed, 33 insertions, 45 deletions

diff --git a/lib/spack/spack/environment/environment.py b/lib/spack/spack/environment/environment.py
index 1353b76b56..44d6785456 100644
--- a/lib/spack/spack/environment/environment.py
+++ b/lib/spack/spack/environment/environment.py
@@ -34,6 +34,7 @@ import spack.spec
 import spack.stage
 import spack.store
 import spack.subprocess_context
+import spack.traverse
 import spack.user_environment as uenv
 import spack.util.cpus
 import spack.util.environment
@@ -1760,22 +1761,14 @@ class Environment(object):
 
     def all_specs(self):
         """Return all specs, even those a user spec would shadow."""
-        all_specs = set()
-        for h in self.concretized_order:
-            try:
-                spec = self.specs_by_hash[h]
-            except KeyError:
-                tty.warn(
-                    "Environment %s appears to be corrupt: missing spec " '"%s"' % (self.name, h)
-                )
-                continue
-            all_specs.update(spec.traverse())
-
-        return sorted(all_specs)
+        roots = [self.specs_by_hash[h] for h in self.concretized_order]
+        specs = [s for s in spack.traverse.traverse_nodes(roots, lambda s: s.dag_hash())]
+        specs.sort()
+        return specs
 
     def all_hashes(self):
         """Return hashes of all specs."""
-        return list(set(s.dag_hash() for s in self.all_specs()))
+        return [s.dag_hash() for s in self.all_specs()]
 
     def roots(self):
         """Specs explicitly requested by the user *in this environment*.
@@ -1812,11 +1805,10 @@ class Environment(object):
 
     def get_by_hash(self, dag_hash):
         matches = {}
-        for _, root in self.concretized_specs():
-            for spec in root.traverse(root=True):
-                dep_hash = spec.dag_hash()
-                if dep_hash.startswith(dag_hash):
-                    matches[dep_hash] = spec
+        roots = [self.specs_by_hash[h] for h in self.concretized_order]
+        for spec in spack.traverse.traverse_nodes(roots, key=lambda s: s.dag_hash()):
+            if spec.dag_hash().startswith(dag_hash):
+                matches[spec.dag_hash()] = spec
         return list(matches.values())
 
     def get_one_by_hash(self, dag_hash):
@@ -1918,28 +1910,27 @@ class Environment(object):
         If these specs appear under different user_specs, only one copy
         is added to the list returned.
         """
-        spec_list = list()
-
-        for spec_hash in self.concretized_order:
-            spec = self.specs_by_hash[spec_hash]
+        specs = [self.specs_by_hash[h] for h in self.concretized_order]
 
-            specs = spec.traverse(deptype=("link", "run")) if recurse_dependencies else (spec,)
-
-            spec_list.extend(specs)
+        if recurse_dependencies:
+            specs.extend(
+                spack.traverse.traverse_nodes(
+                    specs, root=False, deptype=("link", "run"), key=lambda s: s.dag_hash()
+                )
+            )
 
-        return spec_list
+        return specs
 
     def _to_lockfile_dict(self):
         """Create a dictionary to store a lockfile for this environment."""
         concrete_specs = {}
-        for spec in self.specs_by_hash.values():
-            for s in spec.traverse():
-                dag_hash = s.dag_hash()
-                if dag_hash not in concrete_specs:
-                    spec_dict = s.node_dict_with_hashes(hash=ht.dag_hash)
-                    # Assumes no legacy formats, since this was just created.
-                    spec_dict[ht.dag_hash.name] = s.dag_hash()
-                    concrete_specs[dag_hash] = spec_dict
+        for s in spack.traverse.traverse_nodes(
+            self.specs_by_hash.values(), key=lambda s: s.dag_hash()
+        ):
+            spec_dict = s.node_dict_with_hashes(hash=ht.dag_hash)
+            # Assumes no legacy formats, since this was just created.
+            spec_dict[ht.dag_hash.name] = s.dag_hash()
+            concrete_specs[s.dag_hash()] = spec_dict
 
         hash_spec_list = zip(self.concretized_order, self.concretized_user_specs)
 
@@ -2050,19 +2041,16 @@ class Environment(object):
             # ensure the prefix/.env directory exists
             fs.mkdirp(self.env_subdir_path)
 
-            for spec in self.new_specs:
-                for dep in spec.traverse():
-                    if not dep.concrete:
-                        raise ValueError(
-                            "specs passed to environment.write() " "must be concrete!"
-                        )
+            for spec in spack.traverse.traverse_nodes(self.new_specs):
+                if not spec.concrete:
+                    raise ValueError("specs passed to environment.write() " "must be concrete!")
 
-                    root = os.path.join(self.repos_path, dep.namespace)
-                    repo = spack.repo.create_or_construct(root, dep.namespace)
-                    pkg_dir = repo.dirname_for_package_name(dep.name)
+                root = os.path.join(self.repos_path, spec.namespace)
+                repo = spack.repo.create_or_construct(root, spec.namespace)
+                pkg_dir = repo.dirname_for_package_name(spec.name)
 
-                    fs.mkdirp(pkg_dir)
-                    spack.repo.path.dump_provenance(dep, pkg_dir)
+                fs.mkdirp(pkg_dir)
+                spack.repo.path.dump_provenance(spec, pkg_dir)
 
             self._update_and_write_manifest(raw_yaml_dict, yaml_dict)