import codecs
import logging
import re
from os import path
from os import remove
from subprocess import check_call
from tempfile import NamedTemporaryFile
import six
from anytree import PreOrderIter
_RE_ESC = re.compile(r'["\\]')
[docs]class DotExporter(object):
def __init__(self, node, graph="digraph", name="tree", options=None,
indent=4, nodenamefunc=None, nodeattrfunc=None,
edgeattrfunc=None, edgetypefunc=None):
"""
Dot Language Exporter.
Args:
node (Node): start node.
Keyword Args:
graph: DOT graph type.
name: DOT graph name.
options: list of options added to the graph.
indent (int): number of spaces for indent.
nodenamefunc: Function to extract node name from `node` object.
The function shall accept one `node` object as
argument and return the name of it.
nodeattrfunc: Function to decorate a node with attributes.
The function shall accept one `node` object as
argument and return the attributes.
edgeattrfunc: Function to decorate a edge with attributes.
The function shall accept two `node` objects as
argument. The first the node and the second the child
and return the attributes.
edgetypefunc: Function to which gives the edge type.
The function shall accept two `node` objects as
argument. The first the node and the second the child
and return the edge (i.e. '->').
>>> from anytree import Node
>>> root = Node("root")
>>> s0 = Node("sub0", parent=root, edge=2)
>>> s0b = Node("sub0B", parent=s0, foo=4, edge=109)
>>> s0a = Node("sub0A", parent=s0, edge="")
>>> s1 = Node("sub1", parent=root, edge="")
>>> s1a = Node("sub1A", parent=s1, edge=7)
>>> s1b = Node("sub1B", parent=s1, edge=8)
>>> s1c = Node("sub1C", parent=s1, edge=22)
>>> s1ca = Node("sub1Ca", parent=s1c, edge=42)
.. note:: If the node names are not unqiue, see :any:`UniqueDotExporter`.
A directed graph:
>>> from anytree.exporter import DotExporter
>>> for line in DotExporter(root):
... print(line)
digraph tree {
"root";
"sub0";
"sub0B";
"sub0A";
"sub1";
"sub1A";
"sub1B";
"sub1C";
"sub1Ca";
"root" -> "sub0";
"root" -> "sub1";
"sub0" -> "sub0B";
"sub0" -> "sub0A";
"sub1" -> "sub1A";
"sub1" -> "sub1B";
"sub1" -> "sub1C";
"sub1C" -> "sub1Ca";
}
The resulting graph:
.. image:: ../static/dotexporter0.png
An undirected graph:
>>> def nodenamefunc(node):
... return '%s:%s' % (node.name, node.depth)
>>> def edgeattrfunc(node, child):
... return 'label="%s:%s"' % (node.name, child.name)
>>> def edgetypefunc(node, child):
... return '--'
>>> from anytree.exporter import DotExporter
>>> for line in DotExporter(root, graph="graph",
... nodenamefunc=nodenamefunc,
... nodeattrfunc=lambda node: "shape=box",
... edgeattrfunc=edgeattrfunc,
... edgetypefunc=edgetypefunc):
... print(line)
graph tree {
"root:0" [shape=box];
"sub0:1" [shape=box];
"sub0B:2" [shape=box];
"sub0A:2" [shape=box];
"sub1:1" [shape=box];
"sub1A:2" [shape=box];
"sub1B:2" [shape=box];
"sub1C:2" [shape=box];
"sub1Ca:3" [shape=box];
"root:0" -- "sub0:1" [label="root:sub0"];
"root:0" -- "sub1:1" [label="root:sub1"];
"sub0:1" -- "sub0B:2" [label="sub0:sub0B"];
"sub0:1" -- "sub0A:2" [label="sub0:sub0A"];
"sub1:1" -- "sub1A:2" [label="sub1:sub1A"];
"sub1:1" -- "sub1B:2" [label="sub1:sub1B"];
"sub1:1" -- "sub1C:2" [label="sub1:sub1C"];
"sub1C:2" -- "sub1Ca:3" [label="sub1C:sub1Ca"];
}
The resulting graph:
.. image:: ../static/dotexporter1.png
To export custom node implementations or :any:`AnyNode`, please provide a proper `nodenamefunc`:
>>> from anytree import AnyNode
>>> root = AnyNode(id="root")
>>> s0 = AnyNode(id="sub0", parent=root)
>>> s0b = AnyNode(id="s0b", parent=s0)
>>> s0a = AnyNode(id="s0a", parent=s0)
>>> from anytree.exporter import DotExporter
>>> for line in DotExporter(root, nodenamefunc=lambda n: n.id):
... print(line)
digraph tree {
"root";
"sub0";
"s0b";
"s0a";
"root" -> "sub0";
"sub0" -> "s0b";
"sub0" -> "s0a";
}
"""
self.node = node
self.graph = graph
self.name = name
self.options = options
self.indent = indent
self.nodenamefunc = nodenamefunc
self.nodeattrfunc = nodeattrfunc
self.edgeattrfunc = edgeattrfunc
self.edgetypefunc = edgetypefunc
def __iter__(self):
# prepare
indent = " " * self.indent
nodenamefunc = self.nodenamefunc or self._default_nodenamefunc
nodeattrfunc = self.nodeattrfunc or self._default_nodeattrfunc
edgeattrfunc = self.edgeattrfunc or self._default_edgeattrfunc
edgetypefunc = self.edgetypefunc or self._default_edgetypefunc
return self.__iter(indent, nodenamefunc, nodeattrfunc, edgeattrfunc,
edgetypefunc)
@staticmethod
def _default_nodenamefunc(node):
return node.name
@staticmethod
def _default_nodeattrfunc(node):
return None
@staticmethod
def _default_edgeattrfunc(node, child):
return None
@staticmethod
def _default_edgetypefunc(node, child):
return "->"
def __iter(self, indent, nodenamefunc, nodeattrfunc, edgeattrfunc, edgetypefunc):
yield "{self.graph} {self.name} {{".format(self=self)
for option in self.__iter_options(indent):
yield option
for node in self.__iter_nodes(indent, nodenamefunc, nodeattrfunc):
yield node
for edge in self.__iter_edges(indent, nodenamefunc, edgeattrfunc, edgetypefunc):
yield edge
yield "}"
def __iter_options(self, indent):
options = self.options
if options:
for option in options:
yield "%s%s" % (indent, option)
def __iter_nodes(self, indent, nodenamefunc, nodeattrfunc):
for node in PreOrderIter(self.node):
nodename = nodenamefunc(node)
nodeattr = nodeattrfunc(node)
nodeattr = " [%s]" % nodeattr if nodeattr is not None else ""
yield '%s"%s"%s;' % (indent, DotExporter.esc(nodename), nodeattr)
def __iter_edges(self, indent, nodenamefunc, edgeattrfunc, edgetypefunc):
for node in PreOrderIter(self.node):
nodename = nodenamefunc(node)
for child in node.children:
childname = nodenamefunc(child)
edgeattr = edgeattrfunc(node, child)
edgetype = edgetypefunc(node, child)
edgeattr = " [%s]" % edgeattr if edgeattr is not None else ""
yield '%s"%s" %s "%s"%s;' % (indent, DotExporter.esc(nodename), edgetype,
DotExporter.esc(childname), edgeattr)
[docs] def to_dotfile(self, filename):
"""
Write graph to `filename`.
>>> from anytree import Node
>>> root = Node("root")
>>> s0 = Node("sub0", parent=root)
>>> s0b = Node("sub0B", parent=s0)
>>> s0a = Node("sub0A", parent=s0)
>>> s1 = Node("sub1", parent=root)
>>> s1a = Node("sub1A", parent=s1)
>>> s1b = Node("sub1B", parent=s1)
>>> s1c = Node("sub1C", parent=s1)
>>> s1ca = Node("sub1Ca", parent=s1c)
>>> from anytree.exporter import DotExporter
>>> DotExporter(root).to_dotfile("tree.dot")
The generated file should be handed over to the `dot` tool from the
http://www.graphviz.org/ package::
$ dot tree.dot -T png -o tree.png
"""
with codecs.open(filename, "w", "utf-8") as file:
for line in self:
file.write("%s\n" % line)
[docs] def to_picture(self, filename):
"""
Write graph to a temporary file and invoke `dot`.
The output file type is automatically detected from the file suffix.
*`graphviz` needs to be installed, before usage of this method.*
"""
fileformat = path.splitext(filename)[1][1:]
with NamedTemporaryFile("wb", delete=False) as dotfile:
dotfilename = dotfile.name
for line in self:
dotfile.write(("%s\n" % line).encode("utf-8"))
dotfile.flush()
cmd = ["dot", dotfilename, "-T", fileformat, "-o", filename]
check_call(cmd)
try:
remove(dotfilename)
except Exception: # pragma: no cover
msg = 'Could not remove temporary file %s' % dotfilename
logging.getLogger(__name__).warn(msg)
[docs] @staticmethod
def esc(value):
"""Escape Strings."""
return _RE_ESC.sub(lambda m: r"\%s" % m.group(0), six.text_type(value))
[docs]class UniqueDotExporter(DotExporter):
def __init__(self, node, graph="digraph", name="tree", options=None,
indent=4, nodenamefunc=None, nodeattrfunc=None,
edgeattrfunc=None, edgetypefunc=None):
"""
Unqiue Dot Language Exporter.
Handle trees with random or conflicting node names gracefully.
Args:
node (Node): start node.
Keyword Args:
graph: DOT graph type.
name: DOT graph name.
options: list of options added to the graph.
indent (int): number of spaces for indent.
nodenamefunc: Function to extract node name from `node` object.
The function shall accept one `node` object as
argument and return the name of it.
nodeattrfunc: Function to decorate a node with attributes.
The function shall accept one `node` object as
argument and return the attributes.
edgeattrfunc: Function to decorate a edge with attributes.
The function shall accept two `node` objects as
argument. The first the node and the second the child
and return the attributes.
edgetypefunc: Function to which gives the edge type.
The function shall accept two `node` objects as
argument. The first the node and the second the child
and return the edge (i.e. '->').
>>> from anytree import Node
>>> root = Node("root")
>>> s0 = Node("sub0", parent=root)
>>> s0b = Node("s0", parent=s0)
>>> s0a = Node("s0", parent=s0)
>>> s1 = Node("sub1", parent=root)
>>> s1a = Node("s1", parent=s1)
>>> s1b = Node("s1", parent=s1)
>>> s1c = Node("s1", parent=s1)
>>> s1ca = Node("sub1Ca", parent=s1c)
>>> from anytree.exporter import UniqueDotExporter
>>> for line in UniqueDotExporter(root): # doctest: +SKIP
... print(line)
digraph tree {
"0x7f1bf2c9c510" [label="root"];
"0x7f1bf2c9c5a0" [label="sub0"];
"0x7f1bf2c9c630" [label="s0"];
"0x7f1bf2c9c6c0" [label="s0"];
"0x7f1bf2c9c750" [label="sub1"];
"0x7f1bf2c9c7e0" [label="s1"];
"0x7f1bf2c9c870" [label="s1"];
"0x7f1bf2c9c900" [label="s1"];
"0x7f1bf2c9c990" [label="sub1Ca"];
"0x7f1bf2c9c510" -> "0x7f1bf2c9c5a0";
"0x7f1bf2c9c510" -> "0x7f1bf2c9c750";
"0x7f1bf2c9c5a0" -> "0x7f1bf2c9c630";
"0x7f1bf2c9c5a0" -> "0x7f1bf2c9c6c0";
"0x7f1bf2c9c750" -> "0x7f1bf2c9c7e0";
"0x7f1bf2c9c750" -> "0x7f1bf2c9c870";
"0x7f1bf2c9c750" -> "0x7f1bf2c9c900";
"0x7f1bf2c9c900" -> "0x7f1bf2c9c990";
}
The resulting graph:
.. image:: ../static/uniquedotexporter2.png
To export custom node implementations or :any:`AnyNode`, please provide a proper `nodeattrfunc`:
>>> from anytree import AnyNode
>>> root = AnyNode(id="root")
>>> s0 = AnyNode(id="sub0", parent=root)
>>> s0b = AnyNode(id="s0", parent=s0)
>>> s0a = AnyNode(id="s0", parent=s0)
>>> from anytree.exporter import UniqueDotExporter
>>> for line in UniqueDotExporter(root, nodeattrfunc=lambda n: 'label="%s"' % (n.id)): # doctest: +SKIP
... print(line)
digraph tree {
"0x7f5c70449af8" [label="root"];
"0x7f5c70449bd0" [label="sub0"];
"0x7f5c70449c60" [label="s0"];
"0x7f5c70449cf0" [label="s0"];
"0x7f5c70449af8" -> "0x7f5c70449bd0";
"0x7f5c70449bd0" -> "0x7f5c70449c60";
"0x7f5c70449bd0" -> "0x7f5c70449cf0";
}
"""
super(UniqueDotExporter, self).__init__(node, graph=graph, name=name, options=options, indent=indent,
nodenamefunc=nodenamefunc, nodeattrfunc=nodeattrfunc,
edgeattrfunc=edgeattrfunc, edgetypefunc=edgetypefunc)
@staticmethod
def _default_nodenamefunc(node):
return hex(id(node))
@staticmethod
def _default_nodeattrfunc(node):
return 'label="%s"' % (node.name)