ydb/contrib/python/Automat/py3/automat/_visualize.py

from __future__ import annotations

import argparse
import sys
from functools import wraps
from typing import Callable, Iterator

import graphviz

from ._core import Automaton, Input, Output, State
from ._discover import findMachines
from ._methodical import MethodicalMachine
from ._typed import TypeMachine, InputProtocol, Core


def _gvquote(s: str) -> str:
    return '"{}"'.format(s.replace('"', r"\""))


def _gvhtml(s: str) -> str:
    return "<{}>".format(s)


def elementMaker(name: str, *children: str, **attrs: str) -> str:
    """
    Construct a string from the HTML element description.
    """
    formattedAttrs = " ".join(
        "{}={}".format(key, _gvquote(str(value)))
        for key, value in sorted(attrs.items())
    )
    formattedChildren = "".join(children)
    return "<{name} {attrs}>{children}</{name}>".format(
        name=name, attrs=formattedAttrs, children=formattedChildren
    )


def tableMaker(
    inputLabel: str,
    outputLabels: list[str],
    port: str,
    _E: Callable[..., str] = elementMaker,
) -> str:
    """
    Construct an HTML table to label a state transition.
    """
    colspan = {}
    if outputLabels:
        colspan["colspan"] = str(len(outputLabels))

    inputLabelCell = _E(
        "td",
        _E("font", inputLabel, face="menlo-italic"),
        color="purple",
        port=port,
        **colspan,
    )

    pointSize = {"point-size": "9"}
    outputLabelCells = [
        _E("td", _E("font", outputLabel, **pointSize), color="pink")
        for outputLabel in outputLabels
    ]

    rows = [_E("tr", inputLabelCell)]

    if outputLabels:
        rows.append(_E("tr", *outputLabelCells))

    return _E("table", *rows)


def escapify(x: Callable[[State], str]) -> Callable[[State], str]:
    @wraps(x)
    def impl(t: State) -> str:
        return x(t).replace("<", "&lt;").replace(">", "&gt;")

    return impl


def makeDigraph(
    automaton: Automaton[State, Input, Output],
    inputAsString: Callable[[Input], str] = repr,
    outputAsString: Callable[[Output], str] = repr,
    stateAsString: Callable[[State], str] = repr,
) -> graphviz.Digraph:
    """
    Produce a L{graphviz.Digraph} object from an automaton.
    """

    inputAsString = escapify(inputAsString)
    outputAsString = escapify(outputAsString)
    stateAsString = escapify(stateAsString)

    digraph = graphviz.Digraph(
        graph_attr={"pack": "true", "dpi": "100"},
        node_attr={"fontname": "Menlo"},
        edge_attr={"fontname": "Menlo"},
    )

    for state in automaton.states():
        if state is automaton.initialState:
            stateShape = "bold"
            fontName = "Menlo-Bold"
        else:
            stateShape = ""
            fontName = "Menlo"
        digraph.node(
            stateAsString(state),
            fontame=fontName,
            shape="ellipse",
            style=stateShape,
            color="blue",
        )
    for n, eachTransition in enumerate(automaton.allTransitions()):
        inState, inputSymbol, outState, outputSymbols = eachTransition
        thisTransition = "t{}".format(n)
        inputLabel = inputAsString(inputSymbol)

        port = "tableport"
        table = tableMaker(
            inputLabel,
            [outputAsString(outputSymbol) for outputSymbol in outputSymbols],
            port=port,
        )

        digraph.node(thisTransition, label=_gvhtml(table), margin="0.2", shape="none")

        digraph.edge(
            stateAsString(inState),
            "{}:{}:w".format(thisTransition, port),
            arrowhead="none",
        )
        digraph.edge("{}:{}:e".format(thisTransition, port), stateAsString(outState))

    return digraph


def tool(
    _progname: str = sys.argv[0],
    _argv: list[str] = sys.argv[1:],
    _syspath: list[str] = sys.path,
    _findMachines: Callable[
        [str],
        Iterator[tuple[str, MethodicalMachine | TypeMachine[InputProtocol, Core]]],
    ] = findMachines,
    _print: Callable[..., None] = print,
) -> None:
    """
    Entry point for command line utility.
    """

    DESCRIPTION = """
    Visualize automat.MethodicalMachines as graphviz graphs.
    """
    EPILOG = """
    You must have the graphviz tool suite installed.  Please visit
    http://www.graphviz.org for more information.
    """
    if _syspath[0]:
        _syspath.insert(0, "")
    argumentParser = argparse.ArgumentParser(
        prog=_progname, description=DESCRIPTION, epilog=EPILOG
    )
    argumentParser.add_argument(
        "fqpn",
        help="A Fully Qualified Path name" " representing where to find machines.",
    )
    argumentParser.add_argument(
        "--quiet", "-q", help="suppress output", default=False, action="store_true"
    )
    argumentParser.add_argument(
        "--dot-directory",
        "-d",
        help="Where to write out .dot files.",
        default=".automat_visualize",
    )
    argumentParser.add_argument(
        "--image-directory",
        "-i",
        help="Where to write out image files.",
        default=".automat_visualize",
    )
    argumentParser.add_argument(
        "--image-type",
        "-t",
        help="The image format.",
        choices=graphviz.FORMATS,
        default="png",
    )
    argumentParser.add_argument(
        "--view",
        "-v",
        help="View rendered graphs with" " default image viewer",
        default=False,
        action="store_true",
    )
    args = argumentParser.parse_args(_argv)

    explicitlySaveDot = args.dot_directory and (
        not args.image_directory or args.image_directory != args.dot_directory
    )
    if args.quiet:

        def _print(*args):
            pass

    for fqpn, machine in _findMachines(args.fqpn):
        _print(fqpn, "...discovered")

        digraph = machine.asDigraph()

        if explicitlySaveDot:
            digraph.save(filename="{}.dot".format(fqpn), directory=args.dot_directory)
            _print(fqpn, "...wrote dot into", args.dot_directory)

        if args.image_directory:
            deleteDot = not args.dot_directory or explicitlySaveDot
            digraph.format = args.image_type
            digraph.render(
                filename="{}.dot".format(fqpn),
                directory=args.image_directory,
                view=args.view,
                cleanup=deleteDot,
            )
            if deleteDot:
                msg = "...wrote image into"
            else:
                msg = "...wrote image and dot into"
            _print(fqpn, msg, args.image_directory)