From 2271249338d3124d8000a9dcee3ea712ee7133a0 Mon Sep 17 00:00:00 2001 From: Benjy Weinberger Date: Tue, 3 Jan 2023 07:55:01 +0200 Subject: [PATCH] Export a lockfile in an ad-hoc JSON graph format. --- pex/cli/commands/lock.py | 78 +++++++++++++++++++++++------------ pex/resolve/locked_resolve.py | 29 +++++++++++++ 2 files changed, 80 insertions(+), 27 deletions(-) diff --git a/pex/cli/commands/lock.py b/pex/cli/commands/lock.py index 680a9978b..bc35f5bbe 100644 --- a/pex/cli/commands/lock.py +++ b/pex/cli/commands/lock.py @@ -3,9 +3,10 @@ from __future__ import absolute_import, print_function +import json import sys from argparse import Action, ArgumentError, ArgumentParser, ArgumentTypeError, _ActionsContainer -from collections import OrderedDict +from collections import OrderedDict, defaultdict from operator import attrgetter from pex import pex_warnings @@ -65,6 +66,7 @@ class Value(Enum.Value): pass PIP = Value("pip") + GRAPH = Value("graph") PEP_665 = Value("pep-665") @@ -237,8 +239,12 @@ def _add_export_arguments(cls, export_parser): choices=ExportFormat.values(), type=ExportFormat.for_value, help=( - "The format to export the lock to. Currently only the {pip!r} requirements file " - "format using `--hash` is supported.".format(pip=ExportFormat.PIP) + "The format to export the lock to. Export results may be a subset of the full " + "lockfile, targeting a specific interpreter and platform. Currently only " + "`{pip!r}` - a Pip requirements file format using `--hash`, and `{graph!r} - " + "an ad-hoc adjacency graph format, are supported.".format( + pip=ExportFormat.PIP, graph=ExportFormat.GRAPH + ) ), ) export_parser.add_argument( @@ -492,9 +498,11 @@ def dump_with_terminating_newline(out): def _export(self): # type: () -> Result - if self.options.format != ExportFormat.PIP: + if self.options.format not in [ExportFormat.PIP, ExportFormat.GRAPH]: return Error( - "Only the {pip!r} lock format is supported currently.".format(pip=ExportFormat.PIP) + "Only the `{pip!r}` and `{graph!r}` formats are supported currently.".format( + pip=ExportFormat.PIP, graph=ExportFormat.GRAPH + ) ) lockfile_path, lock_file = self._load_lockfile() @@ -502,10 +510,10 @@ def _export(self): resolved_targets = targets.unique_targets() if len(resolved_targets) > 1: return Error( - "A lock can only be exported for a single target in the {pip!r} format.\n" + "A lock can only be exported for a single target in the {export_fmt!r} format.\n" "There were {count} targets selected:\n" "{targets}".format( - pip=ExportFormat.PIP, + export_fmt=self.options.format, count=len(resolved_targets), targets="\n".join( "{index}. {target}".format(index=index, target=target) @@ -535,13 +543,13 @@ def _export(self): resolved_subset.resolved for resolved_subset in subset_result.subsets ) pex_warnings.warn( - "Only a single lock can be exported in the {pip!r} format.\n" + "Only a single lock can be exported in the {export_fmt!r} format.\n" "There were {count} locks stored in {lockfile} that were applicable for the " "selected target: {target}; so using the most specific lock with platform " "{platform}.".format( count=len(subset_result.subsets), lockfile=lockfile_path, - pip=ExportFormat.PIP, + export_fmt=self.options.format, target=target, platform=resolved.source.platform_tag, ) @@ -555,25 +563,41 @@ def _export(self): downloaded_artifact.artifact.fingerprint ) - with self.output(self.options) as output: - pins = fingerprints_by_pin.keys() # type: Iterable[Pin] - if self.options.sort_by == ExportSortBy.PROJECT_NAME: - pins = sorted(pins, key=attrgetter("project_name.normalized")) - for pin in pins: - fingerprints = fingerprints_by_pin[pin] - output.write( - "{project_name}=={version} \\\n" - " {hashes}\n".format( - project_name=pin.project_name, - version=pin.version.raw, - hashes=" \\\n ".join( - "--hash={algorithm}:{hash}".format( - algorithm=fingerprint.algorithm, hash=fingerprint.hash - ) - for fingerprint in fingerprints - ), + if self.options.format == ExportFormat.PIP: + with self.output(self.options) as output: + pins = fingerprints_by_pin.keys() # type: Iterable[Pin] + if self.options.sort_by == ExportSortBy.PROJECT_NAME: + pins = sorted(pins, key=attrgetter("project_name.normalized")) + for pin in pins: + fingerprints = fingerprints_by_pin[pin] + output.write( + "{project_name}=={version} \\\n" + " {hashes}\n".format( + project_name=pin.project_name, + version=pin.version.raw, + hashes=" \\\n ".join( + "--hash={algorithm}:{hash}".format( + algorithm=fingerprint.algorithm, hash=fingerprint.hash + ) + for fingerprint in fingerprints + ), + ) ) - ) + else: + vertices = set() + edges = defaultdict(list) + for src_pin, dst_pins in resolved.adjacency_list.items(): + dst_pin_strs = sorted(str(dst_pin) for dst_pin in dst_pins) + vertices.add(str(src_pin)) + vertices.update(dst_pin_strs) + edges[str(src_pin)] = dst_pin_strs + graph = { + "vertices": sorted(vertices), + "edges": edges, + } + with self.output(self.options) as output: + json.dump(graph, output) + return Ok() def _update(self): diff --git a/pex/resolve/locked_resolve.py b/pex/resolve/locked_resolve.py index 2cf74bd34..eb52ad4c0 100644 --- a/pex/resolve/locked_resolve.py +++ b/pex/resolve/locked_resolve.py @@ -345,6 +345,7 @@ def create( target, # type: Target direct_requirements, # type: Iterable[Requirement] resolved_artifacts, # type: Iterable[_ResolvedArtifact] + adjacency_list, # type: Dict[Pin, Set[Pin]] source, # type: LockedResolve ): # type: (...) -> Resolved @@ -382,6 +383,7 @@ def create( return cls( target_specificity=sum(target_specificities) / len(target_specificities), downloadable_artifacts=tuple(downloadable_artifacts), + adjacency_list=adjacency_list, source=source, ) @@ -396,6 +398,7 @@ def most_specific(cls, resolves): target_specificity = attr.ib() # type: float downloadable_artifacts = attr.ib() # type: Tuple[DownloadableArtifact, ...] + adjacency_list = attr.ib() # type: Dict[Pin, Set[Pin]] source = attr.ib(eq=False) # type: LockedResolve @@ -526,6 +529,7 @@ def resolve( # 1. Gather all required projects and their requirers. required = OrderedDict() # type: OrderedDict[ProjectName, List[_ResolveRequest]] to_be_resolved = deque() # type: Deque[_ResolveRequest] + project_adjacency_list = defaultdict(set) # type Dict[ProjectName, Set[ProjectName]] def request_resolve(requests): # type: (Iterable[_ResolveRequest]) -> None @@ -541,6 +545,14 @@ def request_resolve(requests): resolve_request = to_be_resolved.popleft() project_name = resolve_request.project_name required.setdefault(project_name, []).append(resolve_request) + # Ensure that projects with no requirements appear in the list. + project_adjacency_list[project_name].update([]) + + if len(resolve_request.required_by) > 1: + # NB: resolve_request.required_by[-1] is project_name itself. + project_adjacency_list[resolve_request.required_by[-2].project_name].add( + project_name + ) if not transitive: continue @@ -726,9 +738,26 @@ def attributed_reason(reason): uniqued_resolved_artifacts.append(resolved_artifact) seen.add(resolved_artifact.ranked_artifact.artifact) + # TODO: I assume each project name is unique in uniqued_resolved_artifacts? + project_name_to_pin = { + resolved_artifact.locked_requirement.pin.project_name: resolved_artifact.locked_requirement.pin + for resolved_artifact in uniqued_resolved_artifacts + } + + def _pin(proj_name): + return project_name_to_pin[proj_name] + + pin_adjacency_list = { + _pin(src_project_name): { + _pin(dst_project_name) for dst_project_name in dst_project_names + } + for src_project_name, dst_project_names in project_adjacency_list.items() + } + return Resolved.create( target=target, direct_requirements=requirements, resolved_artifacts=uniqued_resolved_artifacts, + adjacency_list=pin_adjacency_list, source=self, )