zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

traverse: use overload for return type when depth=True vs depth=False (#48482)

Browse Source
This commit is contained in:
Harmen Stoppels 2025-01-10 09:53:28 +01:00 committed by GitHub
parent 464390962f
commit c1d385ada2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 330 additions and 97 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lib/spack/spack/cmd/uninstall.py
View File

@ -144,7 +144,7 @@ def is_installed(spec):
record = spack.store.STORE.db.query_local_by_spec_hash(spec.dag_hash()) record = spack.store.STORE.db.query_local_by_spec_hash(spec.dag_hash())
return record and record.installed return record and record.installed
specs = traverse.traverse_nodes( all_specs = traverse.traverse_nodes(
specs, specs,
root=False, root=False,
order="breadth", order="breadth",
@ -155,7 +155,7 @@ def is_installed(spec):
) )
with spack.store.STORE.db.read_transaction(): with spack.store.STORE.db.read_transaction():
return [spec for spec in specs if is_installed(spec)] return [spec for spec in all_specs if is_installed(spec)]
def dependent_environments( def dependent_environments(

2
lib/spack/spack/database.py
View File

@ -1330,7 +1330,7 @@ def deprecate(self, spec: "spack.spec.Spec", deprecator: "spack.spec.Spec") -> N
def installed_relatives( def installed_relatives(
self, self,
spec: "spack.spec.Spec", spec: "spack.spec.Spec",
direction: str = "children", direction: tr.DirectionType = "children",
transitive: bool = True, transitive: bool = True,
deptype: Union[dt.DepFlag, dt.DepTypes] = dt.ALL, deptype: Union[dt.DepFlag, dt.DepTypes] = dt.ALL,
) -> Set["spack.spec.Spec"]: ) -> Set["spack.spec.Spec"]:

2
lib/spack/spack/installer.py
View File

@ -539,7 +539,7 @@ def dump_packages(spec: "spack.spec.Spec", path: str) -> None:
# Note that we copy them in as they are in the *install* directory # Note that we copy them in as they are in the *install* directory
# NOT as they are in the repository, because we want a snapshot of # NOT as they are in the repository, because we want a snapshot of
# how *this* particular build was done. # how *this* particular build was done.
for node in spec.traverse(deptype=all): for node in spec.traverse(deptype="all"):
if node is not spec: if node is not spec:
# Locate the dependency package in the install tree and find # Locate the dependency package in the install tree and find
# its provenance information. # its provenance information.

147
lib/spack/spack/spec.py
View File

@ -58,7 +58,21 @@
import re import re
import socket import socket
import warnings import warnings
from typing import Any, Callable, Dict, Iterable, List, Match, Optional, Set, Tuple, Union from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Match,
Optional,
Set,
Tuple,
Union,
overload,
)
from typing_extensions import Literal
import archspec.cpu import archspec.cpu
@ -83,7 +97,7 @@
import spack.solver import spack.solver
import spack.spec_parser import spack.spec_parser
import spack.store import spack.store
import spack.traverse as traverse import spack.traverse
import spack.util.executable import spack.util.executable
import spack.util.hash import spack.util.hash
import spack.util.module_cmd as md import spack.util.module_cmd as md
@ -1339,16 +1353,16 @@ def tree(
depth: bool = False, depth: bool = False,
hashes: bool = False, hashes: bool = False,
hashlen: Optional[int] = None, hashlen: Optional[int] = None,
cover: str = "nodes", cover: spack.traverse.CoverType = "nodes",
indent: int = 0, indent: int = 0,
format: str = DEFAULT_FORMAT, format: str = DEFAULT_FORMAT,
deptypes: Union[Tuple[str, ...], str] = "all", deptypes: Union[dt.DepFlag, dt.DepTypes] = dt.ALL,
show_types: bool = False, show_types: bool = False,
depth_first: bool = False, depth_first: bool = False,
recurse_dependencies: bool = True, recurse_dependencies: bool = True,
status_fn: Optional[Callable[["Spec"], InstallStatus]] = None, status_fn: Optional[Callable[["Spec"], InstallStatus]] = None,
prefix: Optional[Callable[["Spec"], str]] = None, prefix: Optional[Callable[["Spec"], str]] = None,
key=id, key: Callable[["Spec"], Any] = id,
) -> str: ) -> str:
"""Prints out specs and their dependencies, tree-formatted with indentation. """Prints out specs and their dependencies, tree-formatted with indentation.
@ -1380,11 +1394,16 @@ def tree(
# reduce deptypes over all in-edges when covering nodes # reduce deptypes over all in-edges when covering nodes
if show_types and cover == "nodes": if show_types and cover == "nodes":
deptype_lookup: Dict[str, dt.DepFlag] = collections.defaultdict(dt.DepFlag) deptype_lookup: Dict[str, dt.DepFlag] = collections.defaultdict(dt.DepFlag)
for edge in traverse.traverse_edges(specs, cover="edges", deptype=deptypes, root=False): for edge in spack.traverse.traverse_edges(
specs, cover="edges", deptype=deptypes, root=False
):
deptype_lookup[edge.spec.dag_hash()] |= edge.depflag deptype_lookup[edge.spec.dag_hash()] |= edge.depflag
for d, dep_spec in traverse.traverse_tree( # SupportsRichComparisonT issue with List[Spec]
sorted(specs), cover=cover, deptype=deptypes, depth_first=depth_first, key=key sorted_specs: List["Spec"] = sorted(specs) # type: ignore[type-var]
for d, dep_spec in spack.traverse.traverse_tree(
sorted_specs, cover=cover, deptype=deptypes, depth_first=depth_first, key=key
): ):
node = dep_spec.spec node = dep_spec.spec
@ -1927,13 +1946,111 @@ def installed_upstream(self):
upstream, _ = spack.store.STORE.db.query_by_spec_hash(self.dag_hash()) upstream, _ = spack.store.STORE.db.query_by_spec_hash(self.dag_hash())
return upstream return upstream
def traverse(self, **kwargs): @overload
"""Shorthand for :meth:`~spack.traverse.traverse_nodes`""" def traverse(
return traverse.traverse_nodes([self], **kwargs) self,
*,
root: bool = ...,
order: spack.traverse.OrderType = ...,
cover: spack.traverse.CoverType = ...,
direction: spack.traverse.DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[False] = False,
key: Callable[["Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable["Spec"]: ...
def traverse_edges(self, **kwargs): @overload
def traverse(
self,
*,
root: bool = ...,
order: spack.traverse.OrderType = ...,
cover: spack.traverse.CoverType = ...,
direction: spack.traverse.DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[True],
key: Callable[["Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[Tuple[int, "Spec"]]: ...
def traverse(
self,
*,
root: bool = True,
order: spack.traverse.OrderType = "pre",
cover: spack.traverse.CoverType = "nodes",
direction: spack.traverse.DirectionType = "children",
deptype: Union[dt.DepFlag, dt.DepTypes] = "all",
depth: bool = False,
key: Callable[["Spec"], Any] = id,
visited: Optional[Set[Any]] = None,
) -> Iterable[Union["Spec", Tuple[int, "Spec"]]]:
"""Shorthand for :meth:`~spack.traverse.traverse_nodes`"""
return spack.traverse.traverse_nodes(
[self],
root=root,
order=order,
cover=cover,
direction=direction,
deptype=deptype,
depth=depth,
key=key,
visited=visited,
)
@overload
def traverse_edges(
self,
*,
root: bool = ...,
order: spack.traverse.OrderType = ...,
cover: spack.traverse.CoverType = ...,
direction: spack.traverse.DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[False] = False,
key: Callable[["Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[DependencySpec]: ...
@overload
def traverse_edges(
self,
*,
root: bool = ...,
order: spack.traverse.OrderType = ...,
cover: spack.traverse.CoverType = ...,
direction: spack.traverse.DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[True],
key: Callable[["Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[Tuple[int, DependencySpec]]: ...
def traverse_edges(
self,
*,
root: bool = True,
order: spack.traverse.OrderType = "pre",
cover: spack.traverse.CoverType = "nodes",
direction: spack.traverse.DirectionType = "children",
deptype: Union[dt.DepFlag, dt.DepTypes] = "all",
depth: bool = False,
key: Callable[["Spec"], Any] = id,
visited: Optional[Set[Any]] = None,
) -> Iterable[Union[DependencySpec, Tuple[int, DependencySpec]]]:
"""Shorthand for :meth:`~spack.traverse.traverse_edges`""" """Shorthand for :meth:`~spack.traverse.traverse_edges`"""
return traverse.traverse_edges([self], **kwargs) return spack.traverse.traverse_edges(
[self],
root=root,
order=order,
cover=cover,
direction=direction,
deptype=deptype,
depth=depth,
key=key,
visited=visited,
)
@property @property
def short_spec(self): def short_spec(self):
@ -4105,10 +4222,10 @@ def tree(
depth: bool = False, depth: bool = False,
hashes: bool = False, hashes: bool = False,
hashlen: Optional[int] = None, hashlen: Optional[int] = None,
cover: str = "nodes", cover: spack.traverse.CoverType = "nodes",
indent: int = 0, indent: int = 0,
format: str = DEFAULT_FORMAT, format: str = DEFAULT_FORMAT,
deptypes: Union[Tuple[str, ...], str] = "all", deptypes: Union[dt.DepTypes, dt.DepFlag] = dt.ALL,
show_types: bool = False, show_types: bool = False,
depth_first: bool = False, depth_first: bool = False,
recurse_dependencies: bool = True, recurse_dependencies: bool = True,

272
lib/spack/spack/traverse.py
View File

@ -3,7 +3,21 @@
# SPDX-License-Identifier: (Apache-2.0 OR MIT) # SPDX-License-Identifier: (Apache-2.0 OR MIT)
from collections import defaultdict from collections import defaultdict
from typing import Any, Callable, List, NamedTuple, Set, Union from typing import (
Any,
Callable,
Iterable,
List,
NamedTuple,
Optional,
Sequence,
Set,
Tuple,
Union,
overload,
)
from typing_extensions import Literal
import spack.deptypes as dt import spack.deptypes as dt
import spack.spec import spack.spec
@ -424,49 +438,95 @@ def traverse_topo_edges_generator(edges, visitor, key=id, root=True, all_edges=F
# High-level API: traverse_edges, traverse_nodes, traverse_tree. # High-level API: traverse_edges, traverse_nodes, traverse_tree.
OrderType = Literal["pre", "post", "breadth", "topo"]
CoverType = Literal["nodes", "edges", "paths"]
DirectionType = Literal["children", "parents"]
@overload
def traverse_edges(
specs: Sequence["spack.spec.Spec"],
*,
root: bool = ...,
order: OrderType = ...,
cover: CoverType = ...,
direction: DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[False] = False,
key: Callable[["spack.spec.Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable["spack.spec.DependencySpec"]: ...
@overload
def traverse_edges(
specs: Sequence["spack.spec.Spec"],
*,
root: bool = ...,
order: OrderType = ...,
cover: CoverType = ...,
direction: DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[True],
key: Callable[["spack.spec.Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[Tuple[int, "spack.spec.DependencySpec"]]: ...
@overload
def traverse_edges(
specs: Sequence["spack.spec.Spec"],
*,
root: bool = ...,
order: OrderType = ...,
cover: CoverType = ...,
direction: DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: bool,
key: Callable[["spack.spec.Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[Union["spack.spec.DependencySpec", Tuple[int, "spack.spec.DependencySpec"]]]: ...
def traverse_edges( def traverse_edges(
specs, specs: Sequence["spack.spec.Spec"],
root=True, root: bool = True,
order="pre", order: OrderType = "pre",
cover="nodes", cover: CoverType = "nodes",
direction="children", direction: DirectionType = "children",
deptype: Union[dt.DepFlag, dt.DepTypes] = "all", deptype: Union[dt.DepFlag, dt.DepTypes] = "all",
depth=False, depth: bool = False,
key=id, key: Callable[["spack.spec.Spec"], Any] = id,
visited=None, visited: Optional[Set[Any]] = None,
): ) -> Iterable[Union["spack.spec.DependencySpec", Tuple[int, "spack.spec.DependencySpec"]]]:
""" """
Generator that yields edges from the DAG, starting from a list of root specs. Iterable of edges from the DAG, starting from a list of root specs.
Arguments: Arguments:
specs (list): List of root specs (considered to be depth 0) specs: List of root specs (considered to be depth 0)
root (bool): Yield the root nodes themselves root: Yield the root nodes themselves
order (str): What order of traversal to use in the DAG. For depth-first order: What order of traversal to use in the DAG. For depth-first search this can be
search this can be ``pre`` or ``post``. For BFS this should be ``breadth``. ``pre`` or ``post``. For BFS this should be ``breadth``. For topological order use
For topological order use ``topo`` ``topo``
cover (str): Determines how extensively to cover the dag. Possible values: cover: Determines how extensively to cover the dag. Possible values:
``nodes`` -- Visit each unique node in the dag only once. ``nodes`` -- Visit each unique node in the dag only once.
``edges`` -- If a node has been visited once but is reached along a ``edges`` -- If a node has been visited once but is reached along a new path, it's
new path, it's accepted, but not recurisvely followed. This traverses accepted, but not recurisvely followed. This traverses each 'edge' in the DAG once.
each 'edge' in the DAG once. ``paths`` -- Explore every unique path reachable from the root. This descends into
``paths`` -- Explore every unique path reachable from the root. visited subtrees and will accept nodes multiple times if they're reachable by multiple
This descends into visited subtrees and will accept nodes multiple paths.
times if they're reachable by multiple paths. direction: ``children`` or ``parents``. If ``children``, does a traversal of this spec's
direction (str): ``children`` or ``parents``. If ``children``, does a traversal children. If ``parents``, traverses upwards in the DAG towards the root.
of this spec's children. If ``parents``, traverses upwards in the DAG
towards the root.
deptype: allowed dependency types deptype: allowed dependency types
depth (bool): When ``False``, yield just edges. When ``True`` yield depth: When ``False``, yield just edges. When ``True`` yield the tuple (depth, edge), where
the tuple (depth, edge), where depth corresponds to the depth depth corresponds to the depth at which edge.spec was discovered.
at which edge.spec was discovered.
key: function that takes a spec and outputs a key for uniqueness test. key: function that takes a spec and outputs a key for uniqueness test.
visited (set or None): a set of nodes not to follow visited: a set of nodes not to follow
Returns: Returns:
A generator that yields ``DependencySpec`` if depth is ``False`` An iterable of ``DependencySpec`` if depth is ``False`` or a tuple of
or a tuple of ``(depth, DependencySpec)`` if depth is ``True``. ``(depth, DependencySpec)`` if depth is ``True``.
""" """
# validate input # validate input
if order == "topo": if order == "topo":
@ -484,7 +544,7 @@ def traverse_edges(
root_edges = with_artificial_edges(specs) root_edges = with_artificial_edges(specs)
# Depth-first # Depth-first
if order in ("pre", "post"): if order == "pre" or order == "post":
return traverse_depth_first_edges_generator( return traverse_depth_first_edges_generator(
root_edges, visitor, order == "post", root, depth root_edges, visitor, order == "post", root, depth
) )
@ -496,79 +556,135 @@ def traverse_edges(
) )
@overload
def traverse_nodes( def traverse_nodes(
specs, specs: Sequence["spack.spec.Spec"],
root=True, *,
order="pre", root: bool = ...,
cover="nodes", order: OrderType = ...,
direction="children", cover: CoverType = ...,
direction: DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[False] = False,
key: Callable[["spack.spec.Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable["spack.spec.Spec"]: ...
@overload
def traverse_nodes(
specs: Sequence["spack.spec.Spec"],
*,
root: bool = ...,
order: OrderType = ...,
cover: CoverType = ...,
direction: DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: Literal[True],
key: Callable[["spack.spec.Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[Tuple[int, "spack.spec.Spec"]]: ...
@overload
def traverse_nodes(
specs: Sequence["spack.spec.Spec"],
*,
root: bool = ...,
order: OrderType = ...,
cover: CoverType = ...,
direction: DirectionType = ...,
deptype: Union[dt.DepFlag, dt.DepTypes] = ...,
depth: bool,
key: Callable[["spack.spec.Spec"], Any] = ...,
visited: Optional[Set[Any]] = ...,
) -> Iterable[Union["spack.spec.Spec", Tuple[int, "spack.spec.Spec"]]]: ...
def traverse_nodes(
specs: Sequence["spack.spec.Spec"],
*,
root: bool = True,
order: OrderType = "pre",
cover: CoverType = "nodes",
direction: DirectionType = "children",
deptype: Union[dt.DepFlag, dt.DepTypes] = "all", deptype: Union[dt.DepFlag, dt.DepTypes] = "all",
depth=False, depth: bool = False,
key=id, key: Callable[["spack.spec.Spec"], Any] = id,
visited=None, visited: Optional[Set[Any]] = None,
): ) -> Iterable[Union["spack.spec.Spec", Tuple[int, "spack.spec.Spec"]]]:
""" """
Generator that yields specs from the DAG, starting from a list of root specs. Iterable of specs from the DAG, starting from a list of root specs.
Arguments: Arguments:
specs (list): List of root specs (considered to be depth 0) specs: List of root specs (considered to be depth 0)
root (bool): Yield the root nodes themselves root: Yield the root nodes themselves
order (str): What order of traversal to use in the DAG. For depth-first order: What order of traversal to use in the DAG. For depth-first search this can be
search this can be ``pre`` or ``post``. For BFS this should be ``breadth``. ``pre`` or ``post``. For BFS this should be ``breadth``.
cover (str): Determines how extensively to cover the dag. Possible values: cover: Determines how extensively to cover the dag. Possible values:
``nodes`` -- Visit each unique node in the dag only once. ``nodes`` -- Visit each unique node in the dag only once.
``edges`` -- If a node has been visited once but is reached along a ``edges`` -- If a node has been visited once but is reached along a new path, it's
new path, it's accepted, but not recurisvely followed. This traverses accepted, but not recurisvely followed. This traverses each 'edge' in the DAG once.
each 'edge' in the DAG once. ``paths`` -- Explore every unique path reachable from the root. This descends into
``paths`` -- Explore every unique path reachable from the root. visited subtrees and will accept nodes multiple times if they're reachable by multiple
This descends into visited subtrees and will accept nodes multiple paths.
times if they're reachable by multiple paths. direction: ``children`` or ``parents``. If ``children``, does a traversal of this spec's
direction (str): ``children`` or ``parents``. If ``children``, does a traversal children. If ``parents``, traverses upwards in the DAG towards the root.
of this spec's children. If ``parents``, traverses upwards in the DAG
towards the root.
deptype: allowed dependency types deptype: allowed dependency types
depth (bool): When ``False``, yield just edges. When ``True`` yield depth: When ``False``, yield just edges. When ``True`` yield the tuple ``(depth, edge)``,
the tuple ``(depth, edge)``, where depth corresponds to the depth where depth corresponds to the depth at which ``edge.spec`` was discovered.
at which ``edge.spec`` was discovered.
key: function that takes a spec and outputs a key for uniqueness test. key: function that takes a spec and outputs a key for uniqueness test.
visited (set or None): a set of nodes not to follow visited: a set of nodes not to follow
Yields: Yields:
By default :class:`~spack.spec.Spec`, or a tuple ``(depth, Spec)`` if depth is By default :class:`~spack.spec.Spec`, or a tuple ``(depth, Spec)`` if depth is
set to ``True``. set to ``True``.
""" """
for item in traverse_edges(specs, root, order, cover, direction, deptype, depth, key, visited): for item in traverse_edges(
yield (item[0], item[1].spec) if depth else item.spec specs,
root=root,
order=order,
cover=cover,
direction=direction,
deptype=deptype,
depth=depth,
key=key,
visited=visited,
):
yield (item[0], item[1].spec) if depth else item.spec # type: ignore
def traverse_tree( def traverse_tree(
specs, cover="nodes", deptype: Union[dt.DepFlag, dt.DepTypes] = "all", key=id, depth_first=True specs: Sequence["spack.spec.Spec"],
): cover: CoverType = "nodes",
deptype: Union[dt.DepFlag, dt.DepTypes] = "all",
key: Callable[["spack.spec.Spec"], Any] = id,
depth_first: bool = True,
) -> Iterable[Tuple[int, "spack.spec.DependencySpec"]]:
""" """
Generator that yields ``(depth, DependencySpec)`` tuples in the depth-first Generator that yields ``(depth, DependencySpec)`` tuples in the depth-first
pre-order, so that a tree can be printed from it. pre-order, so that a tree can be printed from it.
Arguments: Arguments:
specs (list): List of root specs (considered to be depth 0) specs: List of root specs (considered to be depth 0)
cover (str): Determines how extensively to cover the dag. Possible values: cover: Determines how extensively to cover the dag. Possible values:
``nodes`` -- Visit each unique node in the dag only once. ``nodes`` -- Visit each unique node in the dag only once.
``edges`` -- If a node has been visited once but is reached along a ``edges`` -- If a node has been visited once but is reached along a
new path, it's accepted, but not recurisvely followed. This traverses new path, it's accepted, but not recurisvely followed. This traverses each 'edge' in
each 'edge' in the DAG once. the DAG once.
``paths`` -- Explore every unique path reachable from the root. ``paths`` -- Explore every unique path reachable from the root. This descends into
This descends into visited subtrees and will accept nodes multiple visited subtrees and will accept nodes multiple times if they're reachable by multiple
times if they're reachable by multiple paths. paths.
deptype: allowed dependency types deptype: allowed dependency types
key: function that takes a spec and outputs a key for uniqueness test. key: function that takes a spec and outputs a key for uniqueness test.
depth_first (bool): Explore the tree in depth-first or breadth-first order. depth_first: Explore the tree in depth-first or breadth-first order. When setting
When setting ``depth_first=True`` and ``cover=nodes``, each spec only ``depth_first=True`` and ``cover=nodes``, each spec only occurs once at the shallowest
occurs once at the shallowest level, which is useful when rendering level, which is useful when rendering the tree in a terminal.
the tree in a terminal.
Returns: Returns:
A generator that yields ``(depth, DependencySpec)`` tuples in such an order A generator that yields ``(depth, DependencySpec)`` tuples in such an order that a tree can
that a tree can be printed. be printed.
""" """
# BFS only makes sense when going over edges and nodes, for paths the tree is # BFS only makes sense when going over edges and nodes, for paths the tree is
# identical to DFS, which is much more efficient then. # identical to DFS, which is much more efficient then.