diff --git a/src/pgraph/PGraph.py b/src/pgraph/PGraph.py index 8ad82069..2afe42f2 100644 --- a/src/pgraph/PGraph.py +++ b/src/pgraph/PGraph.py @@ -1,6 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +import os import sys import warnings import numpy as np @@ -8,9 +9,8 @@ import copy from collections.abc import Iterable, Iterator import tempfile -import subprocess import webbrowser -from typing import Any, Callable, ClassVar, overload +from typing import IO, Any, Callable, ClassVar, overload from numpy.typing import ArrayLike, NDArray from spatialmath.base.graphics import axes_logic @@ -1444,17 +1444,153 @@ def highlight_vertex( vertex.x, vertex.y, "o", color=color, markersize=12 * scale, alpha=alpha ) - def dotfile(self, filename: str | Any | None = None, direction: str | None = None) -> None: + def export( + self, + filename: str | os.PathLike | IO | None = None, + format: str | None = None, # pylint: disable=redefined-builtin + **options: Any, + ) -> str | None: + """ + Export the graph as text or as a rendered image + + :param filename: destination path or open file, defaults to None + :type filename: str, path-like or file object, optional + :param format: ``"dot"``, ``"mermaid"``, ``"graphml"``, ``"elk"`` or an + image format ``"svg"``, ``"pdf"``, ``"png"`` ..., inferred from + the extension of ``filename`` if not given + :type format: str, optional + :param options: format specific options, see below + :raises ValueError: the format is unknown or cannot be inferred, or an + image format is requested without a ``filename`` + :raises TypeError: an option is not valid for the format + :raises RuntimeError: an image format is requested and GraphViz is not + installed, or fails + :return: the text if ``filename`` is None, otherwise None + :rtype: str or None + + .. list-table:: + :header-rows: 1 + + * - ``format`` + - extensions + - positions + - options + * - ``"dot"`` + - ``.dot`` ``.gv`` + - ``pos="x,y!"`` + - ``direction``, ``colorcomponents``, ``costs`` + * - ``"mermaid"`` + - ``.mmd`` ``.mermaid`` + - no + - ``direction``, ``colorcomponents``, ``costs``, ``fence`` + * - ``"graphml"`` + - ``.graphml`` + - ``x``, ``y``, ``coord`` + - none + * - ``"elk"`` + - ``.elk.json`` ``.elk`` + - ``x``, ``y`` + - ``size``, ``layout_options`` + * - ``"svg"`` ``"pdf"`` ``"png"`` ... + - the same + - yes, via ``neato`` + - ``engine``, ``direction``, ``colorcomponents``, ``costs`` + + The text formats need nothing beyond the standard library. The image + formats run the GraphViz ``dot`` or ``neato`` executable. + + .. runblock:: pycon + + >>> from pgraph import UGraph + >>> g = UGraph() + >>> v1 = g.add_vertex(coord=[0,0], name='v1') + >>> v2 = g.add_vertex(coord=[1,1], name='v2') + >>> g.add_edge(v1, v2, cost=2) + >>> print(g.export(format="mermaid", direction="LR")) + + .. note:: + + - Output is reproducible: vertices are written in creation order + and edges in the order of their vertices, not in set order. + - A file object passed as ``filename`` is written but not closed. + A path is opened, and closed even if writing fails. + - Names are escaped for each format, and Mermaid, GraphML and ELK + use generated ids so any vertex name is valid. + + .. rubric:: For graphs with parallel edges or loops + + Every edge is written, each with its own cost, so parallel edges give + repeated lines (DOT, Mermaid) or distinct ``id`` elements (GraphML, + ELK). A self-loop is written once, as an edge from a vertex to itself. + + :seealso: :meth:`dotfile` :meth:`showgraph` :mod:`pgraph.exporters` + """ + from . import exporters + + if format is None: + format = exporters.format_from_filename(filename) + if format is None: + raise ValueError( + "cannot infer the export format, give format= or a filename " + f"with a known extension; formats: {exporters.formats()}" + ) + format = format.lower() + + if format in exporters.RENDER_FORMATS: + if filename is None: + raise ValueError(f"format {format!r} is an image and needs a filename") + image = exporters.render(self, format, **options) + if isinstance(filename, (str, os.PathLike)): + with open(filename, "wb") as f: + f.write(image) + else: + try: + filename.write(image) + except TypeError: # a text mode file, fine for svg and ps + if not exporters.RENDER_FORMATS[format]: + raise + filename.write(image.decode("utf-8")) + return None + + try: + func = exporters._TEXT[format] + except KeyError: + raise ValueError( + f"unknown export format {format!r}, expected one of " + f"{exporters.formats()}" + ) from None + text = func(self, **options) + + if filename is None: + return text + if isinstance(filename, (str, os.PathLike)): + with open(filename, "w", encoding="utf-8") as f: + f.write(text) + else: + filename.write(text) + return None + + def dotfile( + self, + filename: str | os.PathLike | IO | None = None, + direction: str | None = None, + ) -> None: """ Export graph as a GraphViz dot file :param filename: filename to save graph to, defaults to None - :type filename: str, optional + :type filename: str, path-like or file object, optional + :param direction: layout direction, eg. ``"LR"``, defaults to None + :type direction: str, optional ``g.dotfile()`` creates the specified file which contains the `GraphViz `_ code to represent the embedded graph. By default output is to the console. + This is :meth:`export` with ``format="dot"`` and edge costs left + unlabelled, as it has always been. Use :meth:`export` for other + formats or for edge labels. + .. note:: - The graph is undirected if it is a subclass of ``UGraph`` @@ -1477,76 +1613,44 @@ def dotfile(self, filename: str | Any | None = None, direction: str | None = Non >>> g.add_edge(v1, v2) >>> g.dotfile() - :seealso: :func:`showgraph` + :seealso: :meth:`export` :meth:`showgraph` """ + from . import exporters + text = exporters.dot(self, direction=direction, costs=False) if filename is None: - f = sys.stdout - elif isinstance(filename, str): - f = open(filename, "w") - else: - f = filename - - if isinstance(self, DGraph): - print("digraph {", file=f) + sys.stdout.write(text) + elif isinstance(filename, (str, os.PathLike)): + with open(filename, "w", encoding="utf-8") as f: + f.write(text) else: - print("graph {", file=f) + filename.write(text) - if direction is not None: - print(f"rankdir = {direction}", file=f) - - # add the vertices including name and position - for vertex in self: - if vertex.coord is None: - print(' "{:s}"'.format(vertex.name), file=f) - else: - print( - ' "{:s}" [pos="{:.5g},{:.5g}"]'.format( - vertex.name, vertex.coord[0], vertex.coord[1] - ), - file=f, - ) - print(file=f) - # add the edges - for e in self.edges(): - assert e.v1 is not None and e.v2 is not None - if isinstance(self, DGraph): - print(' "{:s}" -> "{:s}"'.format(e.v1.name, e.v2.name), file=f) - else: - print(' "{:s}" -- "{:s}"'.format(e.v1.name, e.v2.name), file=f) - - print("}", file=f) - - if isinstance(filename, str): - f.close() # noqa - - def showgraph(self, **kwargs: Any) -> None: + def showgraph(self, display: bool = True, **options: Any) -> str | None: """ Display graph in a browser tab - :param kwargs: arguments passed to :meth:`dotfile` - - ``g.showgraph()`` renders and displays the graph in a browser tab. The - graph is exported in `GraphViz `_ format, rendered to - PDF using ``dot`` and then displayed in a browser tab. - - :seealso: :func:`dotfile` - """ - # create the temporary dotfile - dotfile = tempfile.TemporaryFile(mode="w") - self.dotfile(dotfile, **kwargs) - - # rewind the dot file, create PDF file in the filesystem, run dot - dotfile.seek(0) - pdffile = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) - result = subprocess.run("dot -Tpdf", shell=True, stdin=dotfile, stdout=pdffile) - - if result.returncode == 0: - # dot ran happily - # open the PDF file in browser (hopefully portable), then cleanup - webbrowser.open(f"file://{pdffile.name}") - # time.sleep(1) - # os.remove(pdffile.name) + :param display: open the rendered PDF in a browser tab, defaults to True + :param options: ``engine``, ``direction``, ``colorcomponents`` and + ``costs``, passed to :meth:`export` + :raises RuntimeError: GraphViz is not installed, or fails + :return: path of the rendered PDF if ``display`` is False, else None + :rtype: str or None + + ``g.showgraph()`` renders the graph to PDF with `GraphViz + `_ and displays it in a browser tab. With + ``display=False`` the PDF is rendered but not opened and its path is + returned, which is useful for tests and notebooks. + + :seealso: :meth:`export` :meth:`dotfile` + """ + pdf = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) + pdf.close() + self.export(pdf.name, format="pdf", **options) + if display: + webbrowser.open(f"file://{pdf.name}") + return None + return pdf.name def iscyclic(self) -> bool: """ diff --git a/src/pgraph/exporters.py b/src/pgraph/exporters.py new file mode 100644 index 00000000..0ea70e95 --- /dev/null +++ b/src/pgraph/exporters.py @@ -0,0 +1,487 @@ +""" +Export a graph as text (GraphViz DOT, Mermaid, GraphML, ELK JSON) or as a +rendered image (via the GraphViz command line tools). + +This module is reached through :meth:`pgraph.PGraph._BaseGraph.export`, which +dispatches on the format name or the file extension. Each text exporter is a +plain function ``f(graph, **options) -> str``, registered with +:func:`register_exporter`, so a user can add a format without touching the +package:: + + from pgraph.exporters import register_exporter + + @register_exporter("gexf", ".gexf") + def gexf(graph, *, version="1.3") -> str: + ... + +None of the text formats needs a dependency beyond the standard library. Only +the rendered formats (:data:`RENDER_FORMATS`) need the GraphViz executables. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess # nosec B404 - GraphViz is run without a shell, see render() +from typing import IO, TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from .PGraph import BaseVertex, Edge, _BaseGraph + +TextExporter = Callable[..., str] + +_TEXT: dict[str, TextExporter] = {} +_EXTENSIONS: dict[str, str] = {} + +#: GraphViz layout programs that may be run to render an image +ENGINES = ("dot", "neato", "fdp", "sfdp", "twopi", "circo", "osage", "patchwork") + +#: image formats produced by running the GraphViz ``dot``/``neato`` executable, +#: mapped to whether the result is text (``True``) or binary (``False``) +RENDER_FORMATS: dict[str, bool] = { + "svg": True, + "ps": True, + "pdf": False, + "png": False, + "jpg": False, + "jpeg": False, + "gif": False, +} + + +def register_exporter( + name: str, *extensions: str +) -> Callable[[TextExporter], TextExporter]: + """ + Register a text exporter (decorator) + + :param name: format name used as ``export(format=name)`` + :param extensions: file extensions, with the leading dot, that select this + format when ``format`` is not given + :return: the decorator + + The decorated function is called as ``f(graph, **options)`` and returns the + text. Declare its options as keyword-only parameters so that an unknown + option is a ``TypeError``. + """ + + def decorator(func: TextExporter) -> TextExporter: + _TEXT[name] = func + for ext in extensions: + _EXTENSIONS[ext.lower()] = name + return func + + return decorator + + +def formats() -> list[str]: + """ + Names of all supported export formats + + :return: text formats followed by rendered image formats + :rtype: list of str + """ + return sorted(_TEXT) + sorted(RENDER_FORMATS) + + +def format_from_filename(filename: str | os.PathLike | IO | None) -> str | None: + """ + Infer an export format from a file name's extension + + :param filename: path, or a file object with a ``name`` attribute + :return: the format name, or None if it cannot be decided + :rtype: str or None + + Compound extensions such as ``.elk.json`` are tried before the plain + suffix. + """ + if isinstance(filename, (str, os.PathLike)): + name = os.fspath(filename) + else: + name = getattr(filename, "name", None) + if not isinstance(name, str): + return None + lower = name.lower() + for ext in sorted(_EXTENSIONS, key=len, reverse=True): + if lower.endswith(ext): + return _EXTENSIONS[ext] + suffix = os.path.splitext(lower)[1].lstrip(".") + return suffix if suffix in RENDER_FORMATS else None + + +# --------------------------------------------------------------------------- # +# helpers shared by the exporters + + +def _is_directed(graph: _BaseGraph) -> bool: + from .PGraph import DGraph # not at module level: PGraph imports this module + + return isinstance(graph, DGraph) + + +def _ordered_edges(graph: _BaseGraph) -> list[tuple[Edge, BaseVertex, BaseVertex]]: + """ + Edges in a reproducible order + + ``graph.edges()`` is a set, so its order changes from run to run and the + exported text would not be stable (useless for diffs and golden files). + Instead walk the vertices in creation order and each vertex's own edge + list, keeping the first sight of every edge. Identity is by object, not by + endpoints, so parallel edges are all kept, and an undirected edge (held by + both ends, twice for a self-loop) is listed once. + + Each item is ``(edge, start, end)``. + """ + seen: set[int] = set() + edges: list[tuple[Edge, BaseVertex, BaseVertex]] = [] + for vertex in graph: + for e in vertex._edgelist: + if id(e) in seen or e.v1 is None or e.v2 is None: + continue + seen.add(id(e)) + edges.append((e, e.v1, e.v2)) + return edges + + +def _component_colors(graph: _BaseGraph) -> dict[BaseVertex, str]: + """ + Map each vertex to a ``#rrggbb`` colour for its component, as plot() does + """ + import matplotlib.pyplot as plt + import numpy as np + + nc = graph.nc # refreshes vertex.label + cmap = plt.cm.coolwarm(np.linspace(0, 1, max(nc, 1))) + colors = {} + for vertex in graph: + if vertex.label is not None: + r, g, b, _ = cmap[vertex.label] + colors[vertex] = "#{:02x}{:02x}{:02x}".format( + int(r * 255), int(g * 255), int(b * 255) + ) + return colors + + +def _xml_text(value: Any) -> str: + """ + Escape a value for XML element content or a double-quoted attribute + + Characters that XML 1.0 cannot represent at all (most control characters) + are replaced with U+FFFD, since escaping cannot make them valid. + """ + text = str(value) + text = "".join( + c if c in "\t\n\r" or ord(c) >= 0x20 else "\ufffd" for c in text + ) + return ( + text.replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace('"', """) + ) + + +def _dot_quote(text: str) -> str: + """Quote a string for use as a DOT identifier or attribute value""" + escaped = text.replace("\\", "\\\\").replace('"', '\\"') + return '"' + escaped.replace("\r\n", "\\n").replace("\n", "\\n") + '"' + + +# --------------------------------------------------------------------------- # +# text exporters + + +@register_exporter("dot", ".dot", ".gv") +def dot( + graph: _BaseGraph, + *, + direction: str | None = None, + colorcomponents: bool = False, + costs: bool = True, +) -> str: + """ + GraphViz DOT text + + :param direction: layout direction, ``rankdir`` in DOT, eg. ``"LR"`` + :param colorcomponents: fill vertices by connected component + :param costs: label edges with their cost + """ + directed = _is_directed(graph) + link = "->" if directed else "--" + colors = _component_colors(graph) if colorcomponents else {} + + lines = ["digraph {" if directed else "graph {"] + if direction is not None: + lines.append(f" rankdir={_dot_quote(direction)}") + + for vertex in graph: + attrs = [] + if vertex.coord is not None and len(vertex.coord) >= 2: + # the "!" pins the position when the neato engine is used + attrs.append( + 'pos="{:g},{:g}!"'.format(vertex.coord[0], vertex.coord[1]) + ) + if vertex in colors: + attrs.append('style="filled"') + attrs.append(f'fillcolor="{colors[vertex]}"') + suffix = f" [{', '.join(attrs)}]" if attrs else "" + lines.append(f" {_dot_quote(str(vertex.name))}{suffix}") + + for e, v1, v2 in _ordered_edges(graph): + attrs = [] + if costs and e.cost is not None: + attrs.append(f'label="{e.cost:.4g}"') + suffix = f" [{', '.join(attrs)}]" if attrs else "" + lines.append( + f" {_dot_quote(str(v1.name))} {link} {_dot_quote(str(v2.name))}{suffix}" + ) + + lines.append("}") + return "\n".join(lines) + "\n" + + +@register_exporter("mermaid", ".mmd", ".mermaid") +def mermaid( + graph: _BaseGraph, + *, + direction: str = "TB", + colorcomponents: bool = False, + costs: bool = True, + fence: bool = False, +) -> str: + """ + Mermaid flowchart text + + :param direction: ``"TB"``, ``"BT"``, ``"LR"`` or ``"RL"`` + :param colorcomponents: fill vertices by connected component + :param costs: label edges with their cost + :param fence: wrap in a Markdown ```` ```mermaid ```` code fence + + Mermaid flowcharts are laid out automatically and cannot pin node + positions, so vertex coordinates are not represented. Vertices are given + generated ids (``n0``, ``n1``, ...) with the real name as the label, since + names are arbitrary strings but Mermaid ids are not. + """ + if direction not in ("TB", "TD", "BT", "LR", "RL"): + raise ValueError(f"direction must be TB, BT, LR or RL, got {direction!r}") + + arrow = "-->" if _is_directed(graph) else "---" + ids: dict[BaseVertex, str] = {} + lines = [f"flowchart {direction}"] + for i, vertex in enumerate(graph): + ids[vertex] = f"n{i}" + label = str(vertex.name).replace('"', "#quot;") + lines.append(f' n{i}["{label}"]') + + for e, v1, v2 in _ordered_edges(graph): + if costs and e.cost is not None: + lines.append(f' {ids[v1]} {arrow}|"{e.cost:.4g}"| {ids[v2]}') + else: + lines.append(f" {ids[v1]} {arrow} {ids[v2]}") + + if colorcomponents: + colors = _component_colors(graph) + for vertex, color in colors.items(): + lines.append(f" style {ids[vertex]} fill:{color}") + + text = "\n".join(lines) + "\n" + if fence: + text = f"```mermaid\n{text}```\n" + return text + + +@register_exporter("graphml", ".graphml") +def graphml(graph: _BaseGraph) -> str: + """ + GraphML XML text + + Name, component, coordinate (``x``, ``y`` and the full ``coord`` string for + higher dimensions) and edge cost are written as ```` elements with + ```` declarations. ``edgedefault`` follows the graph kind. GraphML + has no standard node position outside the proprietary yFiles extension, + which is not used. + """ + graph.nc # refresh vertex.label + has_coords = any(v.coord is not None for v in graph) + edges = _ordered_edges(graph) + has_cost = any(e.cost is not None for e, _, _ in edges) + + def key(key_id: str, domain: str, name: str, type_: str) -> str: + return ( + f' ' + ) + + kind = "directed" if _is_directed(graph) else "undirected" + lines = [ + '', + '', + key("node_label", "node", "label", "string"), + key("component", "node", "component", "int"), + ] + if has_coords: + lines += [ + key("x", "node", "x", "double"), + key("y", "node", "y", "double"), + key("coord", "node", "coord", "string"), + ] + if has_cost: + lines.append(key("edge_cost", "edge", "cost", "double")) + lines.append(f' ') + + ids: dict[BaseVertex, str] = {} + for i, vertex in enumerate(graph): + ids[vertex] = f"n{i}" + lines.append(f' ') + lines.append(f' {_xml_text(vertex.name)}') + if vertex.label is not None: + lines.append(f' {vertex.label}') + if vertex.coord is not None: + if len(vertex.coord) >= 2: + lines.append(f' {vertex.coord[0]:.6g}') + lines.append(f' {vertex.coord[1]:.6g}') + coord = ",".join(f"{c:.6g}" for c in vertex.coord) + lines.append(f' {coord}') + lines.append(" ") + + for i, (e, v1, v2) in enumerate(edges): + lines.append(f' ') + if e.cost is not None: + lines.append(f' {e.cost:.6g}') + lines.append(" ") + + lines += [" ", ""] + return "\n".join(lines) + "\n" + + +@register_exporter("elk", ".elk.json", ".elk") +def elk( + graph: _BaseGraph, + *, + size: tuple[float, float] = (60, 40), + layout_options: dict[str, str] | None = None, +) -> str: + """ + Eclipse Layout Kernel (ELK) JSON graph text + + :param size: width and height given to every node, which ELK requires + :param layout_options: ELK layout options, merged over the defaults + + Unlike GraphML, ELK has native ``x`` and ``y`` node fields. If any vertex + has a coordinate the ``fixed`` algorithm is selected so that a consumer + keeps the positions; otherwise ``layered`` computes a layout. + """ + graph.nc # refresh vertex.label + has_coords = any(v.coord is not None for v in graph) + + ids: dict[BaseVertex, str] = {} + children = [] + for i, vertex in enumerate(graph): + ids[vertex] = f"n{i}" + node: dict[str, Any] = { + "id": f"n{i}", + "labels": [{"text": str(vertex.name)}], + "width": size[0], + "height": size[1], + } + properties: dict[str, Any] = {} + if vertex.coord is not None: + if len(vertex.coord) >= 2: + node["x"] = float(vertex.coord[0]) + node["y"] = float(vertex.coord[1]) + properties["coord"] = [float(c) for c in vertex.coord] + if vertex.label is not None: + properties["component"] = vertex.label + if properties: + node["properties"] = properties + children.append(node) + + edges = [] + for i, (e, v1, v2) in enumerate(_ordered_edges(graph)): + edge: dict[str, Any] = { + "id": f"e{i}", + "sources": [ids[v1]], + "targets": [ids[v2]], + } + if e.cost is not None: + edge["labels"] = [{"text": f"{e.cost:.4g}"}] + edge["properties"] = {"cost": e.cost} + edges.append(edge) + + options = { + "org.eclipse.elk.algorithm": "fixed" if has_coords else "layered", + "org.eclipse.elk.direction": "DOWN" if _is_directed(graph) else "UNDEFINED", + } + options.update(layout_options or {}) + return ( + json.dumps( + { + "id": "root", + "layoutOptions": options, + "children": children, + "edges": edges, + }, + indent=2, + ) + + "\n" + ) + + +# --------------------------------------------------------------------------- # +# rendered images + + +def render( + graph: _BaseGraph, + fmt: str, + *, + engine: str | None = None, + direction: str | None = None, + colorcomponents: bool = False, + costs: bool = True, +) -> bytes: + """ + Render the graph to an image with the GraphViz executables + + :param fmt: one of :data:`RENDER_FORMATS` + :param engine: GraphViz layout program, one of :data:`ENGINES`. Defaults + to ``neato`` if any vertex has a coordinate, since it honours pinned + positions, else ``dot`` + :return: the image + :rtype: bytes + :raises ValueError: ``fmt`` or ``engine`` is not a known name + :raises RuntimeError: the layout program is not installed, or failed + + Both names are checked against fixed lists and the program is run with an + argument list and no shell, so neither can carry a command. + """ + if fmt not in RENDER_FORMATS: + raise ValueError( + f"unknown image format {fmt!r}, expected one of {sorted(RENDER_FORMATS)}" + ) + if engine is None: + engine = "neato" if any(v.coord is not None for v in graph) else "dot" + if engine not in ENGINES: + raise ValueError(f"unknown GraphViz engine {engine!r}, expected one of {ENGINES}") + exe = shutil.which(engine) + if exe is None: + raise RuntimeError( + f"GraphViz layout program {engine!r} was not found on PATH; " + "install GraphViz (https://graphviz.org/download/) to render " + f"{fmt} output, or export the text with format='dot'" + ) + + text = dot(graph, direction=direction, colorcomponents=colorcomponents, costs=costs) + result = subprocess.run( # nosec B603 nosemgrep - allow-listed program and format, no shell + [exe, f"-T{fmt}"], + input=text.encode("utf-8"), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + if result.returncode != 0: + raise RuntimeError( + f"{engine} failed with exit code {result.returncode}: " + f"{result.stderr.decode('utf-8', 'replace').strip()}" + ) + return result.stdout diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 00000000..0b34c1b2 --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,521 @@ +""" +Graph export: DOT, Mermaid, GraphML, ELK JSON, and rendered images. + +Real I/O boundary: the image formats run the GraphViz executables. Those tests +use the real ``dot``/``neato`` and are skipped when GraphViz is not installed; +the error path for a missing program is tested without needing it. +""" + +import io +import json +import pathlib +import shutil +import tempfile +import unittest +import unittest.mock +import xml.etree.ElementTree as ET # nosec B405 nosemgrep - parses our own generated output + + +from pgraph import DGraph, UGraph, exporters + +HAVE_DOT = shutil.which("dot") is not None +NS = {"g": "http://graphml.graphdrawing.org/xmlns"} + + +def line_graph(Graph): + """a(0,0) -- b(1,0) -- c(2,0) with costs 1 and 2""" + g = Graph() + a = g.add_vertex([0, 0], name="a") + b = g.add_vertex([1, 0], name="b") + c = g.add_vertex([2, 0], name="c") + a.connect(b, cost=1) + b.connect(c, cost=2) + return g, a, b, c + + +class TestRegistry(unittest.TestCase): + def test_formats(self): + names = exporters.formats() + for name in ("dot", "mermaid", "graphml", "elk", "svg", "pdf", "png"): + self.assertIn(name, names) + + def test_format_from_filename(self): + f = exporters.format_from_filename + self.assertEqual(f("a.dot"), "dot") + self.assertEqual(f("a.GV"), "dot") + self.assertEqual(f("a.mmd"), "mermaid") + self.assertEqual(f("a.graphml"), "graphml") + self.assertEqual(f("a.elk.json"), "elk") + self.assertEqual(f(pathlib.Path("x/y.svg")), "svg") + self.assertEqual(f("a.pdf"), "pdf") + self.assertIsNone(f("a.unknown")) + self.assertIsNone(f("noextension")) + self.assertIsNone(f(None)) + self.assertIsNone(f(io.StringIO())) + + def test_file_object_with_name(self): + f = io.StringIO() + f.name = "out.mmd" + self.assertEqual(exporters.format_from_filename(f), "mermaid") + + def test_user_registered_exporter(self): + @exporters.register_exporter("names", ".names") + def names(graph, *, sep=",") -> str: + return sep.join(v.name for v in graph) + + try: + g, *_ = line_graph(UGraph) + self.assertEqual(g.export(format="names"), "a,b,c") + self.assertEqual(g.export(format="names", sep="|"), "a|b|c") + self.assertEqual(exporters.format_from_filename("x.names"), "names") + finally: + exporters._TEXT.pop("names", None) + exporters._EXTENSIONS.pop(".names", None) + + +class TestExportDispatch(unittest.TestCase): + def setUp(self): + self.g, *_ = line_graph(UGraph) + + def test_returns_text_without_filename(self): + text = self.g.export(format="dot") + self.assertIsInstance(text, str) + self.assertTrue(text.startswith("graph {")) + + def test_format_is_case_insensitive(self): + self.assertEqual(self.g.export(format="DOT"), self.g.export(format="dot")) + + def test_cannot_infer(self): + with self.assertRaises(ValueError) as cm: + self.g.export() + self.assertIn("format", str(cm.exception)) + with self.assertRaises(ValueError): + self.g.export("noextension") + + def test_unknown_format(self): + with self.assertRaises(ValueError): + self.g.export(format="bogus") + + def test_unknown_option(self): + with self.assertRaises(TypeError): + self.g.export(format="dot", nonsense=1) + with self.assertRaises(TypeError): + self.g.export(format="graphml", direction="LR") + + def test_image_needs_filename(self): + with self.assertRaises(ValueError): + self.g.export(format="svg") + + def test_path_written_and_returns_none(self): + with tempfile.TemporaryDirectory() as d: + p = pathlib.Path(d) / "g.mmd" + self.assertIsNone(self.g.export(p)) + self.assertEqual(p.read_text(), self.g.export(format="mermaid")) + q = pathlib.Path(d) / "g.txt" + self.assertIsNone(self.g.export(str(q), format="dot")) + self.assertEqual(q.read_text(), self.g.export(format="dot")) + + def test_file_object_not_closed(self): + f = io.StringIO() + self.assertIsNone(self.g.export(f, format="graphml")) + self.assertFalse(f.closed) + self.assertIn(" "b"', text) + self.assertNotIn(" -- ", text) + + def test_options(self): + g, *_ = line_graph(UGraph) + text = g.export(format="dot", direction="LR", costs=False) + self.assertIn('rankdir="LR"', text) + self.assertNotIn("label=", text) + + def test_no_coordinates_no_pos(self): + g = UGraph() + g.add_edge(g.add_vertex(name="p"), g.add_vertex(name="q")) + text = g.export(format="dot") + self.assertNotIn("pos=", text) + self.assertIn('"p" -- "q"', text) + + def test_three_dimensional_uses_first_two(self): + g = UGraph() + g.add_vertex([1, 2, 3], name="p") + self.assertIn('pos="1,2!"', g.export(format="dot")) + + def test_escaping(self): + g = UGraph() + a = g.add_vertex(name='say "hi"') + b = g.add_vertex(name="back\\slash\nnewline") + g.add_edge(a, b) + text = g.export(format="dot") + self.assertIn('"say \\"hi\\""', text) + self.assertIn('"back\\\\slash\\nnewline"', text) + # every line is balanced: no raw newline inside a quoted name + self.assertEqual(len(text.strip().splitlines()), 5) + + def test_colorcomponents(self): + g = UGraph() + a, b, c = (g.add_vertex(name=n) for n in "abc") + g.add_edge(a, b) # c is a second component + text = g.export(format="dot", colorcomponents=True) + colors = { + line.split('fillcolor="')[1].split('"')[0] + for line in text.splitlines() + if "fillcolor" in line + } + self.assertEqual(len(colors), 2) + + +class TestMermaid(unittest.TestCase): + def test_undirected(self): + g, *_ = line_graph(UGraph) + self.assertEqual( + g.export(format="mermaid"), + "flowchart TB\n" + ' n0["a"]\n' + ' n1["b"]\n' + ' n2["c"]\n' + ' n0 ---|"1"| n1\n' + ' n1 ---|"2"| n2\n', + ) + + def test_directed_arrow_and_options(self): + g, *_ = line_graph(DGraph) + text = g.export(format="mermaid", direction="LR", costs=False) + self.assertTrue(text.startswith("flowchart LR")) + self.assertIn("n0 --> n1", text) + self.assertNotIn("|", text) + + def test_fence(self): + g, *_ = line_graph(UGraph) + text = g.export(format="mermaid", fence=True) + self.assertTrue(text.startswith("```mermaid\nflowchart")) + self.assertTrue(text.endswith("```\n")) + + def test_bad_direction(self): + g, *_ = line_graph(UGraph) + with self.assertRaises(ValueError): + g.export(format="mermaid", direction="sideways") + + def test_awkward_names_are_safe(self): + g = UGraph() + g.add_edge(g.add_vertex(name='a "b" c'), g.add_vertex(name="x y-z]")) + text = g.export(format="mermaid") + self.assertIn('n0["a #quot;b#quot; c"]', text) + self.assertIn('n1["x y-z]"]', text) + self.assertIn("n0 --- n1", text) + + def test_colorcomponents(self): + g = UGraph() + g.add_vertex(name="a") + g.add_vertex(name="b") + text = g.export(format="mermaid", colorcomponents=True) + self.assertEqual(sum(1 for l in text.splitlines() if "fill:#" in l), 2) + + +class TestGraphML(unittest.TestCase): + def parse(self, g): + return ET.fromstring(g.export(format="graphml")) # nosec B314 + + def test_structure(self): + g, *_ = line_graph(UGraph) + root = self.parse(g) + graph = root.find("g:graph", NS) + self.assertEqual(graph.get("edgedefault"), "undirected") + self.assertEqual(len(graph.findall("g:node", NS)), 3) + edges = graph.findall("g:edge", NS) + self.assertEqual(len(edges), 2) + self.assertEqual([(e.get("source"), e.get("target")) for e in edges], + [("n0", "n1"), ("n1", "n2")]) + costs = [e.find("g:data", NS).text for e in edges] + self.assertEqual(costs, ["1", "2"]) + + def test_directed(self): + g, *_ = line_graph(DGraph) + graph = self.parse(g).find("g:graph", NS) + self.assertEqual(graph.get("edgedefault"), "directed") + + def test_node_data(self): + g, *_ = line_graph(UGraph) + node = self.parse(g).find("g:graph/g:node[@id='n1']", NS) + data = {d.get("key"): d.text for d in node.findall("g:data", NS)} + self.assertEqual(data["node_label"], "b") + self.assertEqual((data["x"], data["y"]), ("1", "0")) + self.assertEqual(data["coord"], "1,0") + self.assertEqual(data["component"], "0") + + def test_special_characters_roundtrip(self): + g = UGraph() + g.add_vertex(name="<&>\"'") + node = self.parse(g).find("g:graph/g:node", NS) + self.assertEqual(node.find("g:data", NS).text, "<&>\"'") + + def test_unrepresentable_control_characters(self): + # XML 1.0 cannot hold most control characters even escaped; the output + # must still be well formed, with the character replaced + g = UGraph() + g.add_vertex(name="a\x00b\x1fc\td") + node = self.parse(g).find("g:graph/g:node", NS) + self.assertEqual(node.find("g:data", NS).text, "a\ufffdb\ufffdc\td") + + def test_no_cost_key_when_no_costs(self): + g = UGraph() + g.add_edge(g.add_vertex(name="a"), g.add_vertex(name="b")) + text = g.export(format="graphml") + self.assertNotIn("edge_cost", text) + + def test_xml_declaration(self): + g, *_ = line_graph(UGraph) + self.assertTrue( + g.export(format="graphml").startswith('') + ) + + +class TestElk(unittest.TestCase): + def test_structure(self): + g, *_ = line_graph(DGraph) + d = json.loads(g.export(format="elk")) + self.assertEqual(d["id"], "root") + self.assertEqual([c["id"] for c in d["children"]], ["n0", "n1", "n2"]) + self.assertEqual(d["children"][1]["x"], 1.0) + self.assertEqual((d["children"][0]["width"], d["children"][0]["height"]), (60, 40)) + self.assertEqual( + [(e["sources"], e["targets"]) for e in d["edges"]], + [(["n0"], ["n1"]), (["n1"], ["n2"])], + ) + self.assertEqual(d["edges"][1]["properties"]["cost"], 2) + self.assertEqual(d["layoutOptions"]["org.eclipse.elk.algorithm"], "fixed") + self.assertEqual(d["layoutOptions"]["org.eclipse.elk.direction"], "DOWN") + + def test_layered_without_coordinates(self): + g = UGraph() + g.add_edge(g.add_vertex(name="p"), g.add_vertex(name="q")) + d = json.loads(g.export(format="elk")) + self.assertEqual(d["layoutOptions"]["org.eclipse.elk.algorithm"], "layered") + self.assertNotIn("x", d["children"][0]) + + def test_options(self): + g, *_ = line_graph(UGraph) + d = json.loads( + g.export( + format="elk", + size=(10, 5), + layout_options={"org.eclipse.elk.direction": "RIGHT", "k": "v"}, + ) + ) + self.assertEqual(d["children"][0]["width"], 10) + self.assertEqual(d["layoutOptions"]["org.eclipse.elk.direction"], "RIGHT") + self.assertEqual(d["layoutOptions"]["k"], "v") + self.assertEqual(d["layoutOptions"]["org.eclipse.elk.algorithm"], "fixed") + + def test_names_roundtrip(self): + g = UGraph() + g.add_vertex(name='q"uote\\') + d = json.loads(g.export(format="elk")) + self.assertEqual(d["children"][0]["labels"][0]["text"], 'q"uote\\') + + +class TestMultigraph(unittest.TestCase): + """loops and parallel edges are permitted, so every format must cope""" + + def build(self, Graph): + g, a, b, c = line_graph(Graph) + a.connect(b, cost=3) # parallel to a-b + a.connect(b, cost=0.5) # and again + b.connect(b, cost=7) # self-loop + if Graph is DGraph: + b.connect(a, cost=9) # antiparallel + return g + + def expected_edges(self, Graph): + # a-b, b-c, two more a-b, the loop (+ b->a for the directed graph) + return 5 if Graph is UGraph else 6 + + def test_counts(self): + for Graph in (UGraph, DGraph): + with self.subTest(Graph=Graph.__name__): + g = self.build(Graph) + self.assertEqual(g.ne, self.expected_edges(Graph)) + n = self.expected_edges(Graph) + + dot = g.export(format="dot") + link = " -> " if Graph is DGraph else " -- " + self.assertEqual(sum(link in l for l in dot.splitlines()), n) + self.assertIn(f'"b"{link}"b"', dot) + + mm = g.export(format="mermaid") + self.assertEqual(sum("|" in l for l in mm.splitlines()), n) + + gml = ET.fromstring(g.export(format="graphml")) # nosec B314 + edges = gml.findall("g:graph/g:edge", NS) + self.assertEqual(len(edges), n) + self.assertEqual(len({e.get("id") for e in edges}), n) + loops = [e for e in edges if e.get("source") == e.get("target")] + self.assertEqual(len(loops), 1) + + elk = json.loads(g.export(format="elk")) + self.assertEqual(len(elk["edges"]), n) + self.assertEqual(len({e["id"] for e in elk["edges"]}), n) + costs = sorted(e["properties"]["cost"] for e in elk["edges"]) + self.assertIn(0.5, costs) + self.assertIn(3, costs) + + def test_reproducible(self): + # edges() is a set; the export must not depend on its iteration order + for fmt in ("dot", "mermaid", "graphml", "elk"): + with self.subTest(format=fmt): + outputs = {self.build(DGraph).export(format=fmt) for _ in range(5)} + self.assertEqual(len(outputs), 1) + + def test_order_follows_vertices_then_edge_creation(self): + g = self.build(UGraph) + lines = [l for l in g.export(format="dot").splitlines() if "--" in l] + self.assertEqual( + lines, + [ + ' "a" -- "b" [label="1"]', + ' "a" -- "b" [label="3"]', + ' "a" -- "b" [label="0.5"]', + ' "b" -- "c" [label="2"]', # created before the loop + ' "b" -- "b" [label="7"]', + ], + ) + + +class TestWrappers(unittest.TestCase): + def test_dotfile_prints_to_console(self): + g, *_ = line_graph(UGraph) + with unittest.mock.patch("sys.stdout", new_callable=io.StringIO) as out: + g.dotfile() + self.assertEqual(out.getvalue(), g.export(format="dot", costs=False)) + self.assertNotIn("label=", out.getvalue()) + + def test_dotfile_to_path_and_stream(self): + g, *_ = line_graph(UGraph) + with tempfile.TemporaryDirectory() as d: + p = pathlib.Path(d) / "g.dot" + g.dotfile(str(p), direction="LR") + self.assertIn('rankdir="LR"', p.read_text()) + f = io.StringIO() + g.dotfile(f) + self.assertFalse(f.closed) + self.assertTrue(f.getvalue().startswith("graph {")) + + def test_showgraph_missing_graphviz(self): + g, *_ = line_graph(UGraph) + with unittest.mock.patch("shutil.which", return_value=None): + with self.assertRaises(RuntimeError) as cm: + g.showgraph(display=False) + self.assertIn("GraphViz", str(cm.exception)) + + +class TestRender(unittest.TestCase): + def test_engine_not_installed(self): + g, *_ = line_graph(UGraph) + with unittest.mock.patch("shutil.which", return_value=None): + with self.assertRaises(RuntimeError) as cm: + g.export(io.BytesIO(), format="png", engine="neato") + self.assertIn("neato", str(cm.exception)) + + def test_engine_must_be_a_graphviz_program(self): + # the name is run as a program, so it is checked against a fixed list + g, *_ = line_graph(UGraph) + for bad in ("rm", "../dot", "dot; rm -rf /", "no-such-engine"): + with self.subTest(engine=bad): + with self.assertRaises(ValueError): + g.export(io.BytesIO(), format="png", engine=bad) + + def test_unknown_image_format(self): + g, *_ = line_graph(UGraph) + with self.assertRaises(ValueError): + exporters.render(g, "bmp") + + @unittest.skipUnless(HAVE_DOT, "GraphViz not installed") + def test_svg_to_path_and_text_stream(self): + g, *_ = line_graph(UGraph) + with tempfile.TemporaryDirectory() as d: + p = pathlib.Path(d) / "g.svg" + g.export(p) + self.assertIn("]* cx=\"(-?[\d.]+)\"", f.getvalue())] + self.assertEqual(len(xs), 2) + self.assertGreater(abs(xs[1] - xs[0]), 50) + + @unittest.skipUnless(HAVE_DOT, "GraphViz not installed") + def test_multigraph_renders(self): + g, a, b, c = line_graph(UGraph) + a.connect(b, cost=3) + b.connect(b, cost=7) + f = io.BytesIO() + g.export(f, format="png") + self.assertEqual(f.getvalue()[:4], b"\x89PNG") + + @unittest.skipUnless(HAVE_DOT, "GraphViz not installed") + def test_showgraph_display_false_returns_pdf_path(self): + g, *_ = line_graph(UGraph) + with unittest.mock.patch("pgraph.PGraph.webbrowser.open") as opened: + path = g.showgraph(display=False) + opened.assert_not_called() + self.assertEqual(pathlib.Path(path).read_bytes()[:5], b"%PDF-") + + @unittest.skipUnless(HAVE_DOT, "GraphViz not installed") + def test_showgraph_opens_browser(self): + g, *_ = line_graph(UGraph) + with unittest.mock.patch("pgraph.PGraph.webbrowser.open") as opened: + self.assertIsNone(g.showgraph()) + opened.assert_called_once() + self.assertTrue(opened.call_args[0][0].startswith("file://")) + + +if __name__ == "__main__": + unittest.main()