Skip to content

Commit 025c700

Browse files
jsiroisjsirois
authored
Add a dependency graph tool. (#1132)
By default this renders svg which is useful since it includes clickable links to the PyPI page for each node when viewed in a modern browser.
1 parent 43d1130 commit 025c700

File tree

3 files changed

+288
-1
lines changed

3 files changed

+288
-1
lines changed

‎pex/tools/commands/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# Licensed under the Apache License, Version 2.0 (see LICENSE).
33

44
from pex.tools.command import Command
5+
from pex.tools.commands.graph import Graph
56
from pex.tools.commands.info import Info
67
from pex.tools.commands.interpreter import Interpreter
78
from pex.tools.commands.venv import Venv
@@ -13,4 +14,4 @@
1314

1415
def all_commands():
1516
# type: () -> Iterable[Command]
16-
return Info(), Interpreter(), Venv()
17+
return Info(), Interpreter(), Graph(), Venv()

‎pex/tools/commands/digraph.py

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
2+
# Licensed under the Apache License, Version 2.0 (see LICENSE).
3+
4+
from __future__ import absolute_import
5+
6+
from pex.typing import TYPE_CHECKING
7+
8+
if TYPE_CHECKING:
9+
from typing import Dict, IO, List, Mapping, Optional, Tuple
10+
11+
Value = Optional[str]
12+
Attributes = Mapping[str, Value]
13+
14+
15+
class DiGraph(object):
16+
"""Renders a dot digraph built up from nodes and edges."""
17+
18+
@staticmethod
19+
def _render_ID(value):
20+
# type: (str) -> str
21+
# See https://graphviz.org/doc/info/lang.html for the various forms of `ID`.
22+
return '"{}"'.format(value.replace('"', '\\"'))
23+
24+
@classmethod
25+
def _render_a_list(cls, attributes):
26+
# type: (Attributes) -> str
27+
# See https://graphviz.org/doc/info/lang.html for the `a_list` production.
28+
return ", ".join(
29+
"{name}={value}".format(name=name, value=cls._render_ID(value))
30+
for name, value in attributes.items()
31+
if value is not None
32+
)
33+
34+
def __init__(
35+
self,
36+
name, # type: str
37+
strict=True, # type: bool
38+
**attributes # type: Value
39+
):
40+
# type: (...) -> None
41+
"""
42+
:param name: A name for the graph.
43+
:param strict: Whether or not duplicate edges are collapsed into one edge.
44+
"""
45+
self._name = name
46+
self._strict = strict
47+
self._attributes = attributes # type: Attributes
48+
self._nodes = {} # type: Dict[str, Attributes]
49+
self._edges = [] # type: List[Tuple[str, str, Attributes]]
50+
51+
@property
52+
def name(self):
53+
return self._name
54+
55+
def add_node(
56+
self,
57+
name, # type: str
58+
**attributes # type: Value
59+
):
60+
# type: (...) -> None
61+
"""Adds a node to the graph.
62+
63+
This is done implicitly by add_edge for the nodes the edge connects, but may be useful when
64+
the node is either isolated or else needs to be decorated with attributes.
65+
66+
:param name: The name of the node.
67+
"""
68+
self._nodes[name] = attributes
69+
70+
def add_edge(
71+
self,
72+
start, # type: str
73+
end, # type: str
74+
**attributes # type: Value
75+
):
76+
# type: (...) -> None
77+
"""
78+
79+
:param start: The name of the start node.
80+
:param end: The name of the end node.
81+
:param attributes: Any extra attributes for the edge connecting the start node to the end
82+
node.
83+
"""
84+
self._edges.append((start, end, attributes))
85+
86+
def emit(self, out):
87+
# type: (IO[str]) -> None
88+
"""Render the current state of this digraph to the given `out` stream.
89+
90+
:param out: A stream to render this digraph to. N/B.: Will not be flushed or closed.
91+
"""
92+
93+
def emit_attr_stmt(
94+
stmt, # type: str
95+
attributes, # type: Attributes
96+
):
97+
# type: (...) -> None
98+
# See https://graphviz.org/doc/info/lang.html for the `attr_stmt` production.
99+
out.write(
100+
"{statement} [{a_list}];\n".format(
101+
statement=stmt, a_list=self._render_a_list(attributes)
102+
)
103+
)
104+
105+
if self._strict:
106+
out.write("strict ")
107+
out.write("digraph {name} {{\n".format(name=self._render_ID(self._name)))
108+
emit_attr_stmt("graph", self._attributes)
109+
for node, attributes in self._nodes.items():
110+
emit_attr_stmt(self._render_ID(node), attributes)
111+
for start, end, attributes in self._edges:
112+
emit_attr_stmt(
113+
"{start} -> {end}".format(start=self._render_ID(start), end=self._render_ID(end)),
114+
attributes,
115+
)
116+
out.write("}\n")

‎pex/tools/commands/graph.py

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
2+
# Licensed under the Apache License, Version 2.0 (see LICENSE).
3+
4+
from __future__ import absolute_import
5+
6+
import logging
7+
import os
8+
import tempfile
9+
import threading
10+
from argparse import ArgumentParser, Namespace
11+
from contextlib import contextmanager
12+
13+
from pex.common import safe_mkdir
14+
from pex.dist_metadata import requires_dists
15+
from pex.interpreter import PythonInterpreter
16+
from pex.pex import PEX
17+
from pex.tools.command import Command, Ok, OutputMixin, Result, try_open_file, try_run_program
18+
from pex.tools.commands.digraph import DiGraph
19+
from pex.typing import TYPE_CHECKING
20+
from pex.variables import ENV
21+
22+
if TYPE_CHECKING:
23+
from typing import Iterator, IO, Tuple
24+
25+
logger = logging.getLogger(__name__)
26+
27+
28+
class Graph(OutputMixin, Command):
29+
"""Generates a dot graph of the dependencies contained in a PEX file."""
30+
31+
@staticmethod
32+
def _create_dependency_graph(pex):
33+
# type: (PEX) -> DiGraph
34+
graph = DiGraph(
35+
pex.path(),
36+
fontsize="14",
37+
labelloc="t",
38+
label="Dependency graph of {} for interpreter {} ({})".format(
39+
pex.path(), pex.interpreter.binary, pex.interpreter.identity.requirement
40+
),
41+
)
42+
marker_environment = pex.interpreter.identity.env_markers.copy()
43+
marker_environment["extra"] = []
44+
present_dists = frozenset(dist.project_name for dist in pex.activate())
45+
for dist in pex.activate():
46+
graph.add_node(
47+
name=dist.project_name,
48+
label="{name} {version}".format(name=dist.project_name, version=dist.version),
49+
URL="https://pypi.org/project/{name}/{version}".format(
50+
name=dist.project_name, version=dist.version
51+
),
52+
target="_blank",
53+
)
54+
for req in requires_dists(dist):
55+
if (
56+
req.project_name not in present_dists
57+
and req.marker
58+
and not req.marker.evaluate(environment=marker_environment)
59+
):
60+
graph.add_node(
61+
name=req.project_name,
62+
color="lightgrey",
63+
style="filled",
64+
tooltip="inactive requirement",
65+
URL="https://pypi.org/project/{name}".format(name=req.project_name),
66+
target="_blank",
67+
)
68+
graph.add_edge(
69+
start=dist.project_name,
70+
end=req.project_name,
71+
label=str(req) if (req.specifier or req.marker) else None,
72+
fontsize="10",
73+
)
74+
return graph
75+
76+
def add_arguments(self, parser):
77+
# type: (ArgumentParser) -> None
78+
self.add_output_option(parser, entity="dot graph")
79+
parser.add_argument(
80+
"-r",
81+
"--render",
82+
action="store_true",
83+
help="Attempt to render the graph.",
84+
)
85+
parser.add_argument(
86+
"-f",
87+
"--format",
88+
default="svg",
89+
help="The format to render the graph in.",
90+
)
91+
parser.add_argument(
92+
"--open",
93+
action="store_true",
94+
help="Attempt to open the graph in the system viewer (implies --render).",
95+
)
96+
97+
@staticmethod
98+
def _dot(
99+
options, # type: Namespace
100+
graph, # type: DiGraph
101+
render_fp, # type: IO
102+
):
103+
# type: (...) -> Result
104+
read_fd, write_fd = os.pipe()
105+
106+
def emit():
107+
with os.fdopen(write_fd, "w") as fp:
108+
graph.emit(fp)
109+
110+
emit_thread = threading.Thread(name="{} Emitter".format(__name__), target=emit)
111+
emit_thread.daemon = True
112+
emit_thread.start()
113+
114+
try:
115+
return try_run_program(
116+
"dot",
117+
url="https://graphviz.org/",
118+
error="Failed to render dependency graph for {}.".format(graph.name),
119+
args=["-T", options.format],
120+
stdin=read_fd,
121+
stdout=render_fp,
122+
)
123+
finally:
124+
emit_thread.join()
125+
126+
@contextmanager
127+
def _output_for_open(self, options):
128+
# type: (Namespace) -> Iterator[Tuple[IO, str]]
129+
if self.is_stdout(options):
130+
tmpdir = os.path.join(ENV.PEX_ROOT, "tmp")
131+
safe_mkdir(tmpdir)
132+
with tempfile.NamedTemporaryFile(
133+
prefix="{}.".format(__name__),
134+
suffix=".deps.{}".format(options.format),
135+
dir=tmpdir,
136+
delete=False,
137+
) as tmp_out:
138+
yield tmp_out, tmp_out.name
139+
return
140+
141+
with self.output(options, binary=True) as out:
142+
yield out, out.name
143+
144+
def run(
145+
self,
146+
pex, # type: PEX
147+
options, # type: Namespace
148+
):
149+
# type: (...) -> Result
150+
graph = self._create_dependency_graph(pex)
151+
if not (options.render or options.open):
152+
with self.output(options) as out:
153+
graph.emit(out)
154+
return Ok()
155+
156+
if not options.open:
157+
with self.output(options, binary=True) as out:
158+
return self._dot(options, graph, out)
159+
160+
with self._output_for_open(options) as (out, open_path):
161+
result = self._dot(options, graph, out)
162+
if result.is_error:
163+
return result
164+
165+
return try_open_file(
166+
open_path,
167+
error="Failed to open dependency graph of {} rendered in {} for viewing.".format(
168+
pex.path(), open_path
169+
),
170+
)

0 commit comments

Comments
 (0)