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

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

View File

@ -144,7 +144,7 @@ def is_installed(spec):
record = spack.store.STORE.db.query_local_by_spec_hash(spec.dag_hash())
return record and record.installed
specs = traverse.traverse_nodes(
all_specs = traverse.traverse_nodes(
specs,
root=False,
order="breadth",
@ -155,7 +155,7 @@ def is_installed(spec):
)
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(

View File

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

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
# NOT as they are in the repository, because we want a snapshot of
# 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:
# Locate the dependency package in the install tree and find
# its provenance information.

View File

@ -58,7 +58,21 @@
import re
import socket
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
@ -83,7 +97,7 @@
import spack.solver
import spack.spec_parser
import spack.store
import spack.traverse as traverse
import spack.traverse
import spack.util.executable
import spack.util.hash
import spack.util.module_cmd as md
@ -1339,16 +1353,16 @@ def tree(
depth: bool = False,
hashes: bool = False,
hashlen: Optional[int] = None,
cover: str = "nodes",
cover: spack.traverse.CoverType = "nodes",
indent: int = 0,
format: str = DEFAULT_FORMAT,
deptypes: Union[Tuple[str, ...], str] = "all",
deptypes: Union[dt.DepFlag, dt.DepTypes] = dt.ALL,
show_types: bool = False,
depth_first: bool = False,
recurse_dependencies: bool = True,
status_fn: Optional[Callable[["Spec"], InstallStatus]] = None,
prefix: Optional[Callable[["Spec"], str]] = None,
key=id,
key: Callable[["Spec"], Any] = id,
) -> str:
"""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
if show_types and cover == "nodes":
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
for d, dep_spec in traverse.traverse_tree(
sorted(specs), cover=cover, deptype=deptypes, depth_first=depth_first, key=key
# SupportsRichComparisonT issue with List[Spec]
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
@ -1927,13 +1946,111 @@ def installed_upstream(self):
upstream, _ = spack.store.STORE.db.query_by_spec_hash(self.dag_hash())
return upstream
def traverse(self, **kwargs):
"""Shorthand for :meth:`~spack.traverse.traverse_nodes`"""
return traverse.traverse_nodes([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[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`"""
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
def short_spec(self):
@ -4105,10 +4222,10 @@ def tree(
depth: bool = False,
hashes: bool = False,
hashlen: Optional[int] = None,
cover: str = "nodes",
cover: spack.traverse.CoverType = "nodes",
indent: int = 0,
format: str = DEFAULT_FORMAT,
deptypes: Union[Tuple[str, ...], str] = "all",
deptypes: Union[dt.DepTypes, dt.DepFlag] = dt.ALL,
show_types: bool = False,
depth_first: bool = False,
recurse_dependencies: bool = True,

View File

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