From c1d385ada251e09c15f2e7bddc20954ce618f690 Mon Sep 17 00:00:00 2001 From: Harmen Stoppels Date: Fri, 10 Jan 2025 09:53:28 +0100 Subject: [PATCH] traverse: use overload for return type when depth=True vs depth=False (#48482) --- lib/spack/spack/cmd/uninstall.py | 4 +- lib/spack/spack/database.py | 2 +- lib/spack/spack/installer.py | 2 +- lib/spack/spack/spec.py | 147 +++++++++++++++-- lib/spack/spack/traverse.py | 272 ++++++++++++++++++++++--------- 5 files changed, 330 insertions(+), 97 deletions(-) diff --git a/lib/spack/spack/cmd/uninstall.py b/lib/spack/spack/cmd/uninstall.py index f35f133f627..a274497f70e 100644 --- a/lib/spack/spack/cmd/uninstall.py +++ b/lib/spack/spack/cmd/uninstall.py @@ -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( diff --git a/lib/spack/spack/database.py b/lib/spack/spack/database.py index 7ac7328b764..428e1cf62ef 100644 --- a/lib/spack/spack/database.py +++ b/lib/spack/spack/database.py @@ -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"]: diff --git a/lib/spack/spack/installer.py b/lib/spack/spack/installer.py index 316f6b80277..4fc0b734772 100644 --- a/lib/spack/spack/installer.py +++ b/lib/spack/spack/installer.py @@ -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. diff --git a/lib/spack/spack/spec.py b/lib/spack/spack/spec.py index e5b90252b9a..29bd45b4a5c 100644 --- a/lib/spack/spack/spec.py +++ b/lib/spack/spack/spec.py @@ -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, diff --git a/lib/spack/spack/traverse.py b/lib/spack/spack/traverse.py index d9f76072025..7dcf099569a 100644 --- a/lib/spack/spack/traverse.py +++ b/lib/spack/spack/traverse.py @@ -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.