Source code for graphviz.dot

# dot.py - create dot code

r"""Assemble DOT source code objects.

>>> dot = Graph(comment=u'M\xf8nti Pyth\xf8n ik den H\xf8lie Grailen')

>>> dot.node(u'M\xf8\xf8se')
>>> dot.node('trained_by', u'trained by')
>>> dot.node('tutte', u'TUTTE HERMSGERVORDENBROTBORDA')

>>> dot.edge(u'M\xf8\xf8se', 'trained_by')
>>> dot.edge('trained_by', 'tutte')

>>> dot.node_attr['shape'] = 'rectangle'

>>> print(dot.source.replace(u'\xf8', '0'))  #doctest: +NORMALIZE_WHITESPACE
// M0nti Pyth0n ik den H0lie Grailen
graph {
    node [shape=rectangle]
    "M00se"
    trained_by [label="trained by"]
    tutte [label="TUTTE HERMSGERVORDENBROTBORDA"]
    "M00se" -- trained_by
    trained_by -- tutte
}

>>> dot.view('test-output/m00se.gv')  # doctest: +SKIP
'test-output/m00se.gv.pdf'
"""

from . import backend
from . import files
from . import lang

__all__ = ['Graph', 'Digraph']


class Dot(files.File):
    """Assemble, save, and render DOT source code, open result in viewer."""

    _comment = '// %s'
    _subgraph = 'subgraph %s{'
    _subgraph_plain = '%s{'
    _node = _attr = '\t%s%s'
    _attr_plain = _attr % ('%s', '')
    _tail = '}'

    _quote = staticmethod(lang.quote)
    _quote_edge = staticmethod(lang.quote_edge)

    _a_list = staticmethod(lang.a_list)
    _attr_list = staticmethod(lang.attr_list)

    def __init__(self, name=None, comment=None,
                 filename=None, directory=None,
                 format=None, engine=None, encoding=backend.ENCODING,
                 graph_attr=None, node_attr=None, edge_attr=None, body=None,
                 strict=False):
        self.name = name
        self.comment = comment

        super(Dot, self).__init__(filename, directory, format, engine, encoding)

        self.graph_attr = dict(graph_attr) if graph_attr is not None else {}
        self.node_attr = dict(node_attr) if node_attr is not None else {}
        self.edge_attr = dict(edge_attr) if edge_attr is not None else {}

        self.body = list(body) if body is not None else []

        self.strict = strict

    def _kwargs(self):
        result = super(Dot, self)._kwargs()
        result.update(name=self.name,
                      comment=self.comment,
                      graph_attr=dict(self.graph_attr),
                      node_attr=dict(self.node_attr),
                      edge_attr=dict(self.edge_attr),
                      body=list(self.body),
                      strict=self.strict)
        return result

    def clear(self, keep_attrs=False):
        """Reset content to an empty body, clear graph/node/egde_attr mappings.

        Args:
            keep_attrs (bool): preserve graph/node/egde_attr mappings
        """
        if not keep_attrs:
            for a in (self.graph_attr, self.node_attr, self.edge_attr):
                a.clear()
        del self.body[:]

    def __iter__(self, subgraph=False):
        """Yield the DOT source code line by line (as graph or subgraph)."""
        if self.comment:
            yield self._comment % self.comment

        if subgraph:
            if self.strict:
                raise ValueError('subgraphs cannot be strict')
            head = self._subgraph if self.name else self._subgraph_plain
        else:
            head = self._head_strict if self.strict else self._head
        yield head % (self._quote(self.name) + ' ' if self.name else '')

        for kw in ('graph', 'node', 'edge'):
            attrs = getattr(self, '%s_attr' % kw)
            if attrs:
                yield self._attr % (kw, self._attr_list(None, attrs))

        for line in self.body:
            yield line

        yield self._tail

    @property
    def source(self):
        """The DOT source code as string."""
        return '\n'.join(self)

    def node(self, name, label=None, _attributes=None, **attrs):
        """Create a node.

        Args:
            name: Unique identifier for the node inside the source.
            label: Caption to be displayed (defaults to the node ``name``).
            attrs: Any additional node attributes (must be strings).
        """
        name = self._quote(name)
        attr_list = self._attr_list(label, attrs, _attributes)
        line = self._node % (name, attr_list)
        self.body.append(line)

    def edge(self, tail_name, head_name, label=None, _attributes=None, **attrs):
        """Create an edge between two nodes.

        Args:
            tail_name: Start node identifier (format: ``node[:port[:compass]]``).
            head_name: End node identifier (format: ``node[:port[:compass]]``).
            label: Caption to be displayed near the edge.
            attrs: Any additional edge attributes (must be strings).

        Note:
            The ``tail_name`` and ``head_name`` strings are separated by
            (optional) colon(s) into ``node`` name, ``port`` name, and
            ``compass`` (e.g. ``sw``).
            See :ref:`details in the User Guide <ports>`.
        """
        tail_name = self._quote_edge(tail_name)
        head_name = self._quote_edge(head_name)
        attr_list = self._attr_list(label, attrs, _attributes)
        line = self._edge % (tail_name, head_name, attr_list)
        self.body.append(line)

    def edges(self, tail_head_iter):
        """Create a bunch of edges.

        Args:
            tail_head_iter: Iterable of ``(tail_name, head_name)`` pairs
            (format:``node[:port[:compass]]``).


        Note:
            The ``tail_name`` and ``head_name`` strings are separated by
            (optional) colon(s) into ``node`` name, ``port`` name, and
            ``compass`` (e.g. ``sw``).
            See :ref:`details in the User Guide <ports>`.
        """
        edge = self._edge_plain
        quote = self._quote_edge
        lines = (edge % (quote(t), quote(h)) for t, h in tail_head_iter)
        self.body.extend(lines)

    def attr(self, kw=None, _attributes=None, **attrs):
        """Add a general or graph/node/edge attribute statement.

        Args:
            kw: Attributes target (``None`` or ``'graph'``, ``'node'``, ``'edge'``).
            attrs: Attributes to be set (must be strings, may be empty).

        See the :ref:`usage examples in the User Guide <attributes>`.
        """
        if kw is not None and kw.lower() not in ('graph', 'node', 'edge'):
            raise ValueError('attr statement must target graph, node, or edge: '
                '%r' % kw)
        if attrs or _attributes:
            if kw is None:
                a_list = self._a_list(None, attrs, _attributes)
                line = self._attr_plain % a_list
            else:
                attr_list = self._attr_list(None, attrs, _attributes)
                line = self._attr % (kw, attr_list)
            self.body.append(line)

    def subgraph(self, graph=None, name=None, comment=None,
                 graph_attr=None, node_attr=None, edge_attr=None, body=None):
        """Add the current content of the given sole ``graph`` argument as subgraph \
           or return a context manager returning a new graph instance created \
           with the given (``name``, ``comment``, etc.) arguments whose content is \
           added as subgraph when leaving the context manager's ``with``-block.

        Args:
            graph: An instance of the same kind (:class:`.Graph`, :class:`.Digraph`)
                   as the current graph (sole argument in non-with-block use).
            name: Subgraph name (``with``-block use).
            comment: Subgraph comment (``with``-block use).
            graph_attr: Subgraph-level attribute-value mapping (``with``-block use).
            node_attr: Node-level attribute-value mapping (``with``-block use).
            edge_attr: Edge-level attribute-value mapping (``with``-block use).
            body: Verbatim lines to add to the subgraph ``body`` (``with``-block use).

        See the :ref:`usage examples in the User Guide <subgraphs>`.

        When used as a context manager, the returned new graph instance uses
        ``strict=None`` and the parent graph's values for ``directory``,
        ``format``, ``engine``, and ``encoding`` by default.

        Note:
            If the ``name`` of the subgraph begins with ``'cluster'`` (all lowercase)
            the layout engine will treat it as a special cluster subgraph.
        """
        if graph is None:
            return SubgraphContext(self, {'name': name,
                                          'comment': comment,
                                          'directory': self.directory,
                                          'format': self.format,
                                          'engine': self.engine,
                                          'encoding': self.encoding,
                                          'graph_attr': graph_attr,
                                          'node_attr': node_attr,
                                          'edge_attr': edge_attr,
                                          'body': body,
                                          'strict': None})

        args = [name, comment, graph_attr, node_attr, edge_attr, body]
        if not all(a is None for a in args):
            raise ValueError('graph must be sole argument of subgraph()')

        if graph.directed != self.directed:
            raise ValueError('%r cannot add subgraph of different kind:'
                             ' %r' % (self, graph))

        lines = ['\t' + line for line in graph.__iter__(subgraph=True)]
        self.body.extend(lines)


class SubgraphContext(object):
    """Return a blank instance of the parent and add as subgraph on exit."""

    def __init__(self, parent, kwargs):
        self.parent = parent
        self.graph = parent.__class__(**kwargs)

    def __enter__(self):
        return self.graph

    def __exit__(self, type_, value, traceback):
        if type_ is None:
            self.parent.subgraph(self.graph)


[docs]class Graph(Dot): """Graph source code in the DOT language. Args: name: Graph name used in the source code. comment: Comment added to the first line of the source. filename: Filename for saving the source (defaults to ``name`` + ``'.gv'``). directory: (Sub)directory for source saving and rendering. format: Rendering output format (``'pdf'``, ``'png'``, ...). engine: Layout command used (``'dot'``, ``'neato'``, ...). encoding: Encoding for saving the source. graph_attr: Mapping of ``(attribute, value)`` pairs for the graph. node_attr: Mapping of ``(attribute, value)`` pairs set for all nodes. edge_attr: Mapping of ``(attribute, value)`` pairs set for all edges. body: Iterable of verbatim lines to add to the graph ``body``. strict (bool): Rendering should merge multi-edges. Note: All parameters are `optional` and can be changed under their corresponding attribute name after instance creation. """ _head = 'graph %s{' _head_strict = 'strict %s' % _head _edge = '\t%s -- %s%s' _edge_plain = _edge % ('%s', '%s', '') @property def directed(self): """``False``""" return False
[docs]class Digraph(Dot): """Directed graph source code in the DOT language.""" if Graph.__doc__ is not None: __doc__ += Graph.__doc__.partition('.')[2] _head = 'digraph %s{' _head_strict = 'strict %s' % _head _edge = '\t%s -> %s%s' _edge_plain = _edge % ('%s', '%s', '') @property def directed(self): """``True``""" return True