221 lines
7.4 KiB
Python
221 lines
7.4 KiB
Python
|
# This code is licensed under the Apache License, Version 2.0. You may
|
||
|
|
# obtain a copy of this license in the LICENSE.txt file in the root directory
|
||
|
|
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
|
||
|
|
#
|
||
|
|
# Any modifications or derivative works of this code must retain this
|
||
|
|
# copyright notice, and modified files need to carry a notice indicating
|
||
|
|
# that they have been altered from the originals.
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import subprocess
|
||
|
|
import tempfile
|
||
|
|
import io
|
||
|
|
from typing import TypeVar, Callable, cast, TYPE_CHECKING
|
||
|
|
|
||
|
|
from rustworkx import PyDiGraph, PyGraph
|
||
|
|
|
||
|
|
|
||
|
|
try:
|
||
|
|
from PIL import Image # type: ignore
|
||
|
|
|
||
|
|
HAS_PILLOW = True
|
||
|
|
except ImportError:
|
||
|
|
HAS_PILLOW = False
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from PIL import Image # type: ignore
|
||
|
|
|
||
|
|
_S = TypeVar("_S")
|
||
|
|
_T = TypeVar("_T")
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = ["graphviz_draw"]
|
||
|
|
|
||
|
|
METHODS = {"twopi", "neato", "circo", "fdp", "sfdp", "dot"}
|
||
|
|
IMAGE_TYPES = {
|
||
|
|
"canon",
|
||
|
|
"cmap",
|
||
|
|
"cmapx",
|
||
|
|
"cmapx_np",
|
||
|
|
"dia",
|
||
|
|
"dot",
|
||
|
|
"fig",
|
||
|
|
"gd",
|
||
|
|
"gd2",
|
||
|
|
"gif",
|
||
|
|
"hpgl",
|
||
|
|
"imap",
|
||
|
|
"imap_np",
|
||
|
|
"ismap",
|
||
|
|
"jpe",
|
||
|
|
"jpeg",
|
||
|
|
"jpg",
|
||
|
|
"mif",
|
||
|
|
"mp",
|
||
|
|
"pcl",
|
||
|
|
"pdf",
|
||
|
|
"pic",
|
||
|
|
"plain",
|
||
|
|
"plain-ext",
|
||
|
|
"png",
|
||
|
|
"ps",
|
||
|
|
"ps2",
|
||
|
|
"svg",
|
||
|
|
"svgz",
|
||
|
|
"vml",
|
||
|
|
"vmlz",
|
||
|
|
"vrml",
|
||
|
|
"vtx",
|
||
|
|
"wbmp",
|
||
|
|
"xdor",
|
||
|
|
"xlib",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def graphviz_draw(
|
||
|
|
graph: "PyDiGraph[_S, _T] | PyGraph[_S, _T]", # noqa
|
||
|
|
node_attr_fn: Callable[[_S], dict[str, str]] | None = None,
|
||
|
|
edge_attr_fn: Callable[[_T], dict[str, str]] | None = None,
|
||
|
|
graph_attr: dict[str, str] | None = None,
|
||
|
|
filename: str | None = None,
|
||
|
|
image_type: str | None = None,
|
||
|
|
method: str | None = None,
|
||
|
|
) -> Image | None:
|
||
|
|
"""Draw a :class:`~rustworkx.PyGraph` or :class:`~rustworkx.PyDiGraph` object
|
||
|
|
using graphviz
|
||
|
|
|
||
|
|
.. note::
|
||
|
|
|
||
|
|
This requires that pydot, pillow, and graphviz be installed. Pydot can
|
||
|
|
be installed via pip with ``pip install pydot pillow`` however graphviz
|
||
|
|
will need to be installed separately. You can refer to the
|
||
|
|
Graphviz
|
||
|
|
`documentation <https://graphviz.org/download/#executable-packages>`__
|
||
|
|
for instructions on how to install it.
|
||
|
|
|
||
|
|
:param graph: The rustworkx graph object to draw, can be a
|
||
|
|
:class:`~rustworkx.PyGraph` or a :class:`~rustworkx.PyDiGraph`
|
||
|
|
:param node_attr_fn: An optional callable object that will be passed the
|
||
|
|
weight/data payload for every node in the graph and expected to return
|
||
|
|
a dictionary of Graphviz node attributes to be associated with the node
|
||
|
|
in the visualization. The key and value of this dictionary **must** be
|
||
|
|
a string.
|
||
|
|
:param edge_attr_fn: An optional callable that will be passed the
|
||
|
|
weight/data payload for each edge in the graph and expected to return a
|
||
|
|
dictionary of Graphviz edge attributes to be associated with the edge
|
||
|
|
in the visualization file. The key and value of this dictionary
|
||
|
|
must be a string.
|
||
|
|
:param dict graph_attr: An optional dictionary that specifies any Graphviz
|
||
|
|
graph attributes for the visualization. The key and value of this
|
||
|
|
dictionary must be a string.
|
||
|
|
:param str filename: An optional path to write the visualization to. If
|
||
|
|
specified the return type from this function will be ``None`` as the
|
||
|
|
output image is saved to disk.
|
||
|
|
:param str image_type: The image file format to use for the generated
|
||
|
|
visualization. The support image formats are:
|
||
|
|
``'canon'``, ``'cmap'``, ``'cmapx'``, ``'cmapx_np'``, ``'dia'``,
|
||
|
|
``'dot'``, ``'fig'``, ``'gd'``, ``'gd2'``, ``'gif'``, ``'hpgl'``,
|
||
|
|
``'imap'``, ``'imap_np'``, ``'ismap'``, ``'jpe'``, ``'jpeg'``,
|
||
|
|
``'jpg'``, ``'mif'``, ``'mp'``, ``'pcl'``, ``'pdf'``, ``'pic'``,
|
||
|
|
``'plain'``, ``'plain-ext'``, ``'png'``, ``'ps'``, ``'ps2'``,
|
||
|
|
``'svg'``, ``'svgz'``, ``'vml'``, ``'vmlz'``, ``'vrml'``, ``'vtx'``,
|
||
|
|
``'wbmp'``, ``'xdot'``, ``'xlib'``. It's worth noting that while these
|
||
|
|
formats can all be used for generating image files when the ``filename``
|
||
|
|
kwarg is specified, the Pillow library used for the returned object can
|
||
|
|
not work with all these formats.
|
||
|
|
:param str method: The layout method/Graphviz command method to use for
|
||
|
|
generating the visualization. Available options are ``'dot'``,
|
||
|
|
``'twopi'``, ``'neato'``, ``'circo'``, ``'fdp'``, and ``'sfdp'``.
|
||
|
|
You can refer to the
|
||
|
|
`Graphviz documentation <https://graphviz.org/documentation/>`__ for
|
||
|
|
more details on the different layout methods. By default ``'dot'`` is
|
||
|
|
used.
|
||
|
|
|
||
|
|
:returns: A ``PIL.Image`` object of the generated visualization, if
|
||
|
|
``filename`` is not specified. If ``filename`` is specified then
|
||
|
|
``None`` will be returned as the visualization was written to the
|
||
|
|
path specified in ``filename``
|
||
|
|
:rtype: PIL.Image
|
||
|
|
|
||
|
|
.. jupyter-execute::
|
||
|
|
|
||
|
|
import rustworkx as rx
|
||
|
|
from rustworkx.visualization import graphviz_draw
|
||
|
|
|
||
|
|
def node_attr(node):
|
||
|
|
if node == 0:
|
||
|
|
return {'color': 'yellow', 'fillcolor': 'yellow', 'style': 'filled'}
|
||
|
|
if node % 2:
|
||
|
|
return {'color': 'blue', 'fillcolor': 'blue', 'style': 'filled'}
|
||
|
|
else:
|
||
|
|
return {'color': 'red', 'fillcolor': 'red', 'style': 'filled'}
|
||
|
|
|
||
|
|
graph = rx.generators.directed_star_graph(weights=list(range(32)))
|
||
|
|
graphviz_draw(graph, node_attr_fn=node_attr, method='sfdp')
|
||
|
|
|
||
|
|
"""
|
||
|
|
if not HAS_PILLOW:
|
||
|
|
raise ImportError(
|
||
|
|
"Pillow is necessary to use graphviz_draw() "
|
||
|
|
"it can be installed with 'pip install pydot pillow'"
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
subprocess.run(
|
||
|
|
["dot", "-V"],
|
||
|
|
cwd=tempfile.gettempdir(),
|
||
|
|
check=True,
|
||
|
|
capture_output=True,
|
||
|
|
)
|
||
|
|
except Exception:
|
||
|
|
raise RuntimeError(
|
||
|
|
"Graphviz could not be found or run. This function requires that "
|
||
|
|
"Graphviz is installed. If you need to install Graphviz you can "
|
||
|
|
"refer to: https://graphviz.org/download/#executable-packages for "
|
||
|
|
"instructions."
|
||
|
|
)
|
||
|
|
|
||
|
|
dot_str = cast(str, graph.to_dot(node_attr_fn, edge_attr_fn, graph_attr))
|
||
|
|
if image_type is None:
|
||
|
|
output_format = "png"
|
||
|
|
else:
|
||
|
|
if image_type not in IMAGE_TYPES:
|
||
|
|
raise ValueError(
|
||
|
|
"The specified value for the image_type argument, "
|
||
|
|
f"'{image_type}' is not a valid choice. It must be one of: "
|
||
|
|
f"{IMAGE_TYPES}"
|
||
|
|
)
|
||
|
|
output_format = image_type
|
||
|
|
|
||
|
|
if method is None:
|
||
|
|
prog = "dot"
|
||
|
|
else:
|
||
|
|
if method not in METHODS:
|
||
|
|
raise ValueError(
|
||
|
|
f"The specified value for the method argument, '{method}' is "
|
||
|
|
f"not a valid choice. It must be one of: {METHODS}"
|
||
|
|
)
|
||
|
|
prog = method
|
||
|
|
|
||
|
|
if not filename:
|
||
|
|
dot_result = subprocess.run(
|
||
|
|
[prog, "-T", output_format],
|
||
|
|
input=dot_str.encode("utf-8"),
|
||
|
|
capture_output=True,
|
||
|
|
encoding=None,
|
||
|
|
check=True,
|
||
|
|
text=False,
|
||
|
|
)
|
||
|
|
dot_bytes_image = io.BytesIO(dot_result.stdout)
|
||
|
|
image = Image.open(dot_bytes_image)
|
||
|
|
return image
|
||
|
|
else:
|
||
|
|
subprocess.run(
|
||
|
|
[prog, "-T", output_format, "-o", filename],
|
||
|
|
input=dot_str,
|
||
|
|
check=True,
|
||
|
|
encoding="utf8",
|
||
|
|
text=True,
|
||
|
|
)
|
||
|
|
return None
|