graphviz_helper.py

import os
import subprocess
import sys
import tempfile
from typing import List, Optional, Tuple, Union
import numpy as np
from onnx import ModelProto


def _find_in_PATH(prog: str) -> Optional[str]:
    """
    Looks into every path mentioned in ``%PATH%`` a specific file,
    it raises an exception if not found.

    :param prog: program to look for
    :return: path
    """
    sep = ";" if sys.platform.startswith("win") else ":"
    path = os.environ["PATH"]
    for p in path.split(sep):
        f = os.path.join(p, prog)
        if os.path.exists(f):
            return p
    return None


def _find_graphviz_dot(exc: bool = True) -> str:
    """
    Determines the path to graphviz (on Windows),
    the function tests the existence of versions 34 to 45
    assuming it was installed in a standard folder:
    ``C:\\Program Files\\MiKTeX 2.9\\miktex\\bin\\x64``.

    :param exc: raise exception of be silent
    :return: path to dot
    :raises FileNotFoundError: if graphviz not found
    """
    if sys.platform.startswith("win"):
        version = list(range(34, 60))
        version.extend([f"{v}.1" for v in version])
        for v in version:
            graphviz_dot = f"C:\\Program Files (x86)\\Graphviz2.{v}\\bin\\dot.exe"
            if os.path.exists(graphviz_dot):
                return graphviz_dot
        extra = ["build/update_modules/Graphviz/bin"]
        for ext in extra:
            graphviz_dot = os.path.join(ext, "dot.exe")
            if os.path.exists(graphviz_dot):
                return graphviz_dot
        p = _find_in_PATH("dot.exe")
        if p is None:
            if exc:
                raise FileNotFoundError(
                    f"Unable to find graphviz, look into paths such as {graphviz_dot}."
                )
            return None
        return os.path.join(p, "dot.exe")
    # linux
    return "dot"


def _run_subprocess(
    args: List[str],
    cwd: Optional[str] = None,
):
    assert not isinstance(
        args, str
    ), "args should be a sequence of strings, not a string."

    p = subprocess.Popen(
        args,
        cwd=cwd,
        shell=False,
        env=os.environ,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    raise_exception = False
    output = ""
    while True:
        output = p.stdout.readline().decode(errors="ignore")
        if output == "" and p.poll() is not None:
            break
        if output:
            if (
                "fatal error" in output
                or "CMake Error" in output
                or "gmake: ***" in output
                or "): error C" in output
                or ": error: " in output
            ):
                raise_exception = True
    p.poll()
    error = p.stderr.readline().decode(errors="ignore")
    p.stdout.close()
    if error and raise_exception:
        raise RuntimeError(
            f"An error was found in the output. The build is stopped."
            f"\n{output}\n---\n{error}"
        )
    return output + "\n" + error


def _run_graphviz(filename: str, image: str, engine: str = "dot") -> str:
    """
    Run :epkg:`Graphviz`.

    :param filename: filename which contains the graph definition
    :param image: output image
    :param engine: *dot* or *neato*
    :return: output of graphviz
    """
    ext = os.path.splitext(image)[-1]
    assert ext in {
        ".png",
        ".bmp",
        ".fig",
        ".gif",
        ".ico",
        ".jpg",
        ".jpeg",
        ".pdf",
        ".ps",
        ".svg",
        ".vrml",
        ".tif",
        ".tiff",
        ".wbmp",
    }, f"Unexpected extension {ext!r} for {image!r}."
    if sys.platform.startswith("win"):
        bin_ = os.path.dirname(_find_graphviz_dot())
        # if bin not in os.environ["PATH"]:
        #    os.environ["PATH"] = os.environ["PATH"] + ";" + bin
        exe = os.path.join(bin_, engine)
    else:
        exe = engine
    if os.path.exists(image):
        os.remove(image)
    cmd = [exe, f"-T{ext[1:]}", filename, "-o", image]
    output = _run_subprocess(cmd)
    assert os.path.exists(image), (
        f"Unable to find {image!r}, command line is "
        f"{' '.join(cmd)!r}, Graphviz failed due to\n{output}"
    )
    return output


def draw_graph_graphviz(
    dot: Union[str, ModelProto],
    image: str,
    engine: str = "dot",
) -> str:
    """
    Draws a graph using :epkg:`Graphviz`.

    :param dot: dot graph or ModelProto
    :param image: output image, None, just returns the output
    :param engine: *dot* or *neato*
    :return: :epkg:`Graphviz` output or
        the dot text if *image* is None

    The function creates a temporary file to store the dot file if *image* is not None.
    """
    if isinstance(dot, ModelProto):
        from .dot_plot import to_dot

        sdot = to_dot(dot)
    else:
        sdot = dot
    with tempfile.NamedTemporaryFile(delete=False) as fp:
        fp.write(sdot.encode("utf-8"))
        fp.close()

        filename = fp.name
        assert os.path.exists(
            filename
        ), f"File {filename!r} cannot be created to store the graph."
        out = _run_graphviz(filename, image, engine=engine)
        assert os.path.exists(
            image
        ), f"Graphviz failed with no reason, {image!r} not found, output is {out}."
        os.remove(filename)
        return out


def plot_dot(
    dot: Union[str, ModelProto],
    ax: Optional["matplotlib.axis.Axis"] = None,  # noqa: F821
    engine: str = "dot",
    figsize: Optional[Tuple[int, int]] = None,
) -> "matplotlib.axis.Axis":  # noqa: F821
    """
    Draws a dot graph into a matplotlib graph.

    :param dot: dot graph or ModelProto
    :param image: output image, None, just returns the output
    :param engine: *dot* or *neato*
    :param figsize: figsize of ax is None
    :return: :epkg:`Graphviz` output or, the dot text if *image* is None

    .. plot::

        import matplotlib.pyplot as plt
        import onnx.parser
        from onnx_array_api.plotting.graphviz_helper import plot_dot

        model = onnx.parser.parse_model(
            '''
            <ir_version: 8, opset_import: [ "": 18]>
            agraph (float[N] x) => (float[N] z) {
                two = Constant <value_float=2.0> ()
                four = Add(two, two)
                z = Mul(four, four)
            }
        ''')

        ax = plot_dot(model)
        ax.set_title("Dummy graph")
        plt.show()
    """
    if ax is None:
        import matplotlib.pyplot as plt

        _, ax = plt.subplots(1, 1, figsize=figsize)
        clean = True
    else:
        clean = False

    from PIL import Image

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as fp:
        fp.close()

        draw_graph_graphviz(dot, fp.name, engine=engine)
        img = np.asarray(Image.open(fp.name))
        os.remove(fp.name)

        ax.imshow(img)

    if clean:
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        ax.get_figure().tight_layout()
    return ax