Skip to content

Commit

Permalink
Add a dependency graph tool.
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jsirois committed Dec 14, 2020
1 parent 43d1130 commit 66d12e8
Show file tree
Hide file tree
Showing 3 changed files with 278 additions and 1 deletion.
3 changes: 2 additions & 1 deletion pex/tools/commands/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from pex.tools.command import Command
from pex.tools.commands.graph import Graph
from pex.tools.commands.info import Info
from pex.tools.commands.interpreter import Interpreter
from pex.tools.commands.venv import Venv
Expand All @@ -13,4 +14,4 @@

def all_commands():
# type: () -> Iterable[Command]
return Info(), Interpreter(), Venv()
return Info(), Interpreter(), Graph(), Venv()
116 changes: 116 additions & 0 deletions pex/tools/commands/digraph.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import

from pex.typing import TYPE_CHECKING

if TYPE_CHECKING:
from typing import Dict, IO, List, Mapping, Optional, Tuple

Value = Optional[str]
Attributes = Mapping[str, Value]


class DiGraph(object):
"""Renders a dot digraph built up from nodes and edges."""

@staticmethod
def _render_ID(value):
# type: (str) -> str
# See https://graphviz.org/doc/info/lang.html for the various forms of `ID`.
return '"{}"'.format(value.replace('"', '\\"'))

@classmethod
def _render_a_list(cls, attributes):
# type: (Attributes) -> str
# See https://graphviz.org/doc/info/lang.html for the `a_list` production.
return ", ".join(
"{name}={value}".format(name=name, value=cls._render_ID(value))
for name, value in attributes.items()
if value is not None
)

def __init__(
self,
name, # type: str
strict=True, # type: bool
**attributes # type: Value
):
# type: (...) -> None
"""
:param name: A name for the graph.
:param strict: Whether or not duplicate edges are collapsed into one edge.
"""
self._name = name
self._strict = strict
self._attributes = attributes # type: Attributes
self._nodes = {} # type: Dict[str, Attributes]
self._edges = [] # type: List[Tuple[str, str, Attributes]]

@property
def name(self):
return self._name

def add_node(
self,
name, # type: str
**attributes # type: Value
):
# type: (...) -> None
"""Adds a node to the graph.
This is done implicitly by add_edge for the nodes the edge connects, but may be useful when
the node is either isolated or else needs to be decorated with attributes.
:param name: The name of the node.
"""
self._nodes[name] = attributes

def add_edge(
self,
start, # type: str
end, # type: str
**attributes # type: Value
):
# type: (...) -> None
"""
:param start: The name of the start node.
:param end: The name of the end node.
:param attributes: Any extra attributes for the edge connecting the start node to the end
node.
"""
self._edges.append((start, end, attributes))

def emit(self, out):
# type: (IO[str]) -> None
"""Render the current state of this digraph to the given `out` stream.
:param out: A stream to render this digraph to. N/B.: Will not be flushed or closed.
"""

def emit_attr_stmt(
stmt, # type: str
attributes, # type: Attributes
):
# type: (...) -> None
# See https://graphviz.org/doc/info/lang.html for the `attr_stmt` production.
out.write(
"{statement} [{a_list}];\n".format(
statement=stmt, a_list=self._render_a_list(attributes)
)
)

if self._strict:
out.write("strict ")
out.write("digraph {name} {{\n".format(name=self._render_ID(self._name)))
emit_attr_stmt("graph", self._attributes)
for node, attributes in self._nodes.items():
emit_attr_stmt(self._render_ID(node), attributes)
for start, end, attributes in self._edges:
emit_attr_stmt(
"{start} -> {end}".format(start=self._render_ID(start), end=self._render_ID(end)),
attributes,
)
out.write("}\n")
160 changes: 160 additions & 0 deletions pex/tools/commands/graph.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import

import logging
import os
import tempfile
import threading
from argparse import ArgumentParser, Namespace
from contextlib import contextmanager

from pex.common import safe_mkdir
from pex.dist_metadata import requires_dists
from pex.interpreter import PythonInterpreter
from pex.pex import PEX
from pex.tools.command import Command, Ok, OutputMixin, Result, try_open_file, try_run_program
from pex.tools.commands.digraph import DiGraph
from pex.typing import TYPE_CHECKING
from pex.variables import ENV

if TYPE_CHECKING:
from typing import Iterator, IO, Tuple

logger = logging.getLogger(__name__)


class Graph(OutputMixin, Command):
"""Generates a dot graph of the dependencies contained in a PEX file."""

@staticmethod
def _create_dependency_graph(pex):
graph = DiGraph(pex.path(), fontsize="14")
marker_environment = PythonInterpreter.get().identity.env_markers.copy()
marker_environment["extra"] = []
present_dists = frozenset(dist.project_name for dist in pex.activate())
for dist in pex.activate():
graph.add_node(
name=dist.project_name,
label="{name} {version}".format(name=dist.project_name, version=dist.version),
URL="https://pypi.org/project/{name}/{version}".format(
name=dist.project_name, version=dist.version
),
)
for req in requires_dists(dist):
if (
req.project_name not in present_dists
and req.marker
and not req.marker.evaluate(environment=marker_environment)
):
graph.add_node(
name=req.project_name,
color="lightgrey",
style="filled",
tooltip="inactive requirement",
URL="https://pypi.org/project/{name}".format(name=req.project_name),
)
graph.add_edge(
start=dist.project_name,
end=req.project_name,
label=str(req) if (req.specifier or req.marker) else None,
fontsize="10",
)
return graph

def add_arguments(self, parser):
# type: (ArgumentParser) -> None
self.add_output_option(parser, entity="dot graph")
parser.add_argument(
"-r",
"--render",
action="store_true",
help="Attempt to render the graph.",
)
parser.add_argument(
"-f",
"--format",
default="svg",
help="The format to render the graph in.",
)
parser.add_argument(
"--open",
action="store_true",
help="Attempt to open the graph in the system viewer (implies --render).",
)

@staticmethod
def _dot(
options, # type: Namespace
graph, # type: DiGraph
render_fp, # type: IO
):
# type: (...) -> Result
read_fd, write_fd = os.pipe()

def emit():
with os.fdopen(write_fd, "w") as fp:
graph.emit(fp)

emit_thread = threading.Thread(name="{} Emitter".format(__name__), target=emit)
emit_thread.daemon = True
emit_thread.start()

try:
return try_run_program(
"dot",
url="https://graphviz.org/",
error="Failed to render dependency graph for {}.".format(graph.name),
args=["-T", options.format],
stdin=read_fd,
stdout=render_fp,
)
finally:
emit_thread.join()

@contextmanager
def _output_for_open(self, options):
# type: (Namespace) -> Iterator[Tuple[IO, str]]
if self.is_stdout(options):
tmpdir = os.path.join(ENV.PEX_ROOT, "tmp")
safe_mkdir(tmpdir)
with tempfile.NamedTemporaryFile(
prefix="{}.".format(__name__),
suffix=".deps.{}".format(options.format),
dir=tmpdir,
delete=False,
) as tmp_out:
yield tmp_out, tmp_out.name
return

with self.output(options, binary=True) as out:
yield out, out.name

def run(
self,
pex, # type: PEX
options, # type: Namespace
):
# type: (...) -> Result
graph = self._create_dependency_graph(pex)
if not (options.render or options.open):
with self.output(options) as out:
graph.emit(out)
return Ok()

if not options.open:
with self.output(options, binary=True) as out:
return self._dot(options, graph, out)

with self._output_for_open(options) as (out, open_path):
result = self._dot(options, graph, out)
if result.is_error:
return result

return try_open_file(
open_path,
error="Failed to open dependency graph of {} rendered in {} for viewing.".format(
pex.path(), open_path
),
)

0 comments on commit 66d12e8

Please sign in to comment.