Skip to content

Commit b337439

Browse files
jsiroisjsirois
committed
Add a dependency graph tool.
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 85b9eaf commit b337439

File tree

3 files changed

+278
-1
lines changed

3 files changed

+278
-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: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
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+
graph = DiGraph(pex.path(), fontsize="14")
34+
marker_environment = PythonInterpreter.get().identity.env_markers.copy()
35+
marker_environment["extra"] = []
36+
present_dists = frozenset(dist.project_name for dist in pex.activate())
37+
for dist in pex.activate():
38+
graph.add_node(
39+
name=dist.project_name,
40+
label="{name} {version}".format(name=dist.project_name, version=dist.version),
41+
URL="https://pypi.org/project/{name}/{version}".format(
42+
name=dist.project_name, version=dist.version
43+
),
44+
)
45+
for req in requires_dists(dist):
46+
if (
47+
req.project_name not in present_dists
48+
and req.marker
49+
and not req.marker.evaluate(environment=marker_environment)
50+
):
51+
graph.add_node(
52+
name=req.project_name,
53+
color="lightgrey",
54+
style="filled",
55+
tooltip="inactive requirement",
56+
URL="https://pypi.org/project/{name}".format(name=req.project_name),
57+
)
58+
graph.add_edge(
59+
start=dist.project_name,
60+
end=req.project_name,
61+
label=str(req) if (req.specifier or req.marker) else None,
62+
fontsize="10",
63+
)
64+
return graph
65+
66+
def add_arguments(self, parser):
67+
# type: (ArgumentParser) -> None
68+
self.add_output_option(parser, entity="dot graph")
69+
parser.add_argument(
70+
"-r",
71+
"--render",
72+
action="store_true",
73+
help="Attempt to render the graph.",
74+
)
75+
parser.add_argument(
76+
"-f",
77+
"--format",
78+
default="svg",
79+
help="The format to render the graph in.",
80+
)
81+
parser.add_argument(
82+
"--open",
83+
action="store_true",
84+
help="Attempt to open the graph in the system viewer (implies --render).",
85+
)
86+
87+
@staticmethod
88+
def _dot(
89+
options, # type: Namespace
90+
graph, # type: DiGraph
91+
render_fp, # type: IO
92+
):
93+
# type: (...) -> Result
94+
read_fd, write_fd = os.pipe()
95+
96+
def emit():
97+
with os.fdopen(write_fd, "w") as fp:
98+
graph.emit(fp)
99+
100+
emit_thread = threading.Thread(name="{} Emitter".format(__name__), target=emit)
101+
emit_thread.daemon = True
102+
emit_thread.start()
103+
104+
try:
105+
return try_run_program(
106+
"dot",
107+
url="https://graphviz.org/",
108+
error="Failed to render dependency graph for {}.".format(graph.name),
109+
args=["-T", options.format],
110+
stdin=read_fd,
111+
stdout=render_fp,
112+
)
113+
finally:
114+
emit_thread.join()
115+
116+
@contextmanager
117+
def _output_for_open(self, options):
118+
# type: (Namespace) -> Iterator[Tuple[IO, str]]
119+
if self.is_stdout(options):
120+
tmpdir = os.path.join(ENV.PEX_ROOT, "tmp")
121+
safe_mkdir(tmpdir)
122+
with tempfile.NamedTemporaryFile(
123+
prefix="{}.".format(__name__),
124+
suffix=".deps.{}".format(options.format),
125+
dir=tmpdir,
126+
delete=False,
127+
) as tmp_out:
128+
yield tmp_out, tmp_out.name
129+
return
130+
131+
with self.output(options, binary=True) as out:
132+
yield out, out.name
133+
134+
def run(
135+
self,
136+
pex, # type: PEX
137+
options, # type: Namespace
138+
):
139+
# type: (...) -> Result
140+
graph = self._create_dependency_graph(pex)
141+
if not (options.render or options.open):
142+
with self.output(options) as out:
143+
graph.emit(out)
144+
return Ok()
145+
146+
if not options.open:
147+
with self.output(options, binary=True) as out:
148+
return self._dot(options, graph, out)
149+
150+
with self._output_for_open(options) as (out, open_path):
151+
result = self._dot(options, graph, out)
152+
if result.is_error:
153+
return result
154+
155+
return try_open_file(
156+
open_path,
157+
error="Failed to open dependency graph of {} rendered in {} for viewing.".format(
158+
pex.path(), open_path
159+
),
160+
)

0 commit comments

Comments
 (0)