|
|
import codecs |
|
|
import itertools |
|
|
import logging |
|
|
import re |
|
|
from os import path, remove |
|
|
from subprocess import check_call |
|
|
from tempfile import NamedTemporaryFile |
|
|
|
|
|
import six |
|
|
|
|
|
from anytree import PreOrderIter |
|
|
|
|
|
_RE_ESC = re.compile(r'["\\]') |
|
|
|
|
|
|
|
|
class DotExporter: |
|
|
""" |
|
|
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. '->'). |
|
|
|
|
|
filter_: Function to filter nodes to include in export. |
|
|
The function shall accept one `node` object as |
|
|
argument and return True if it should be included, |
|
|
or False if it should not be included. |
|
|
|
|
|
stop: stop iteration at `node` if `stop` function returns `True` for `node`. |
|
|
|
|
|
maxlevel (int): Limit export to this number of levels. |
|
|
|
|
|
>>> 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"; |
|
|
} |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
node, |
|
|
graph="digraph", |
|
|
name="tree", |
|
|
options=None, |
|
|
indent=4, |
|
|
nodenamefunc=None, |
|
|
nodeattrfunc=None, |
|
|
edgeattrfunc=None, |
|
|
edgetypefunc=None, |
|
|
filter_=None, |
|
|
maxlevel=None, |
|
|
stop=None, |
|
|
): |
|
|
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 |
|
|
self.filter_ = filter_ |
|
|
self.maxlevel = maxlevel |
|
|
self.stop = stop |
|
|
|
|
|
def __iter__(self): |
|
|
|
|
|
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 |
|
|
filter_ = self.filter_ or self._default_filter |
|
|
return self.__iter(indent, nodenamefunc, nodeattrfunc, edgeattrfunc, edgetypefunc, filter_) |
|
|
|
|
|
@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 "->" |
|
|
|
|
|
@staticmethod |
|
|
def _default_filter(node): |
|
|
|
|
|
return True |
|
|
|
|
|
def __iter(self, indent, nodenamefunc, nodeattrfunc, edgeattrfunc, edgetypefunc, filter_): |
|
|
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, filter_): |
|
|
yield node |
|
|
for edge in self.__iter_edges(indent, nodenamefunc, edgeattrfunc, edgetypefunc, filter_): |
|
|
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, filter_): |
|
|
for node in PreOrderIter(self.node, filter_=filter_, stop=self.stop, maxlevel=self.maxlevel): |
|
|
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, filter_): |
|
|
maxlevel = self.maxlevel - 1 if self.maxlevel else None |
|
|
for node in PreOrderIter(self.node, filter_=filter_, stop=self.stop, maxlevel=maxlevel): |
|
|
nodename = nodenamefunc(node) |
|
|
for child in node.children: |
|
|
if not filter_(child): |
|
|
continue |
|
|
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, |
|
|
) |
|
|
|
|
|
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) |
|
|
|
|
|
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: |
|
|
logging.getLogger(__name__).warning("Could not remove temporary file %s", dotfilename) |
|
|
|
|
|
@staticmethod |
|
|
def esc(value): |
|
|
"""Escape Strings.""" |
|
|
return _RE_ESC.sub(lambda m: r"\%s" % m.group(0), six.text_type(value)) |
|
|
|
|
|
|
|
|
class UniqueDotExporter(DotExporter): |
|
|
""" |
|
|
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. '->'). |
|
|
|
|
|
filter_: Function to filter nodes to include in export. |
|
|
The function shall accept one `node` object as |
|
|
argument and return True if it should be included, |
|
|
or False if it should not be included. |
|
|
|
|
|
stop: stop iteration at `node` if `stop` function returns `True` for `node`. |
|
|
|
|
|
maxlevel (int): Limit export to this number of levels. |
|
|
|
|
|
>>> 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): |
|
|
... print(line) |
|
|
digraph tree { |
|
|
"0x0" [label="root"]; |
|
|
"0x1" [label="sub0"]; |
|
|
"0x2" [label="s0"]; |
|
|
"0x3" [label="s0"]; |
|
|
"0x4" [label="sub1"]; |
|
|
"0x5" [label="s1"]; |
|
|
"0x6" [label="s1"]; |
|
|
"0x7" [label="s1"]; |
|
|
"0x8" [label="sub1Ca"]; |
|
|
"0x0" -> "0x1"; |
|
|
"0x0" -> "0x4"; |
|
|
"0x1" -> "0x2"; |
|
|
"0x1" -> "0x3"; |
|
|
"0x4" -> "0x5"; |
|
|
"0x4" -> "0x6"; |
|
|
"0x4" -> "0x7"; |
|
|
"0x7" -> "0x8"; |
|
|
} |
|
|
|
|
|
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)): |
|
|
... print(line) |
|
|
digraph tree { |
|
|
"0x0" [label="root"]; |
|
|
"0x1" [label="sub0"]; |
|
|
"0x2" [label="s0"]; |
|
|
"0x3" [label="s0"]; |
|
|
"0x0" -> "0x1"; |
|
|
"0x1" -> "0x2"; |
|
|
"0x1" -> "0x3"; |
|
|
} |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
node, |
|
|
graph="digraph", |
|
|
name="tree", |
|
|
options=None, |
|
|
indent=4, |
|
|
nodenamefunc=None, |
|
|
nodeattrfunc=None, |
|
|
edgeattrfunc=None, |
|
|
edgetypefunc=None, |
|
|
filter_=None, |
|
|
stop=None, |
|
|
maxlevel=None, |
|
|
): |
|
|
super(UniqueDotExporter, self).__init__( |
|
|
node, |
|
|
graph=graph, |
|
|
name=name, |
|
|
options=options, |
|
|
indent=indent, |
|
|
nodenamefunc=nodenamefunc, |
|
|
nodeattrfunc=nodeattrfunc, |
|
|
edgeattrfunc=edgeattrfunc, |
|
|
edgetypefunc=edgetypefunc, |
|
|
filter_=filter_, |
|
|
stop=stop, |
|
|
maxlevel=maxlevel, |
|
|
) |
|
|
self.__node_ids = {} |
|
|
self.__node_counter = itertools.count() |
|
|
|
|
|
|
|
|
def _default_nodenamefunc(self, node): |
|
|
node_id = id(node) |
|
|
try: |
|
|
num = self.__node_ids[node_id] |
|
|
except KeyError: |
|
|
num = self.__node_ids[node_id] = next(self.__node_counter) |
|
|
return hex(num) |
|
|
|
|
|
@staticmethod |
|
|
def _default_nodeattrfunc(node): |
|
|
return 'label="%s"' % (node.name,) |
|
|
|
