123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394 |
- 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'["\\]')
- class DotExporter(object):
- def __init__(self, node, graph="digraph", name="tree", options=None,
- indent=4, nodenamefunc=None, nodeattrfunc=None,
- edgeattrfunc=None, edgetypefunc=None, maxlevel=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. '->').
- 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";
- }
- """
- 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.maxlevel = maxlevel
- 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, 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):
- maxlevel = self.maxlevel - 1 if self.maxlevel else None
- for node in PreOrderIter(self.node, maxlevel=maxlevel):
- 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)
- 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: # pragma: no cover
- msg = 'Could not remove temporary file %s' % dotfilename
- logging.getLogger(__name__).warn(msg)
- @staticmethod
- def esc(value):
- """Escape Strings."""
- return _RE_ESC.sub(lambda m: r"\%s" % m.group(0), six.text_type(value))
- 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)
|