dotexporter.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394
  1. import codecs
  2. import logging
  3. import re
  4. from os import path
  5. from os import remove
  6. from subprocess import check_call
  7. from tempfile import NamedTemporaryFile
  8. import six
  9. from anytree import PreOrderIter
  10. _RE_ESC = re.compile(r'["\\]')
  11. class DotExporter(object):
  12. def __init__(self, node, graph="digraph", name="tree", options=None,
  13. indent=4, nodenamefunc=None, nodeattrfunc=None,
  14. edgeattrfunc=None, edgetypefunc=None, maxlevel=None):
  15. """
  16. Dot Language Exporter.
  17. Args:
  18. node (Node): start node.
  19. Keyword Args:
  20. graph: DOT graph type.
  21. name: DOT graph name.
  22. options: list of options added to the graph.
  23. indent (int): number of spaces for indent.
  24. nodenamefunc: Function to extract node name from `node` object.
  25. The function shall accept one `node` object as
  26. argument and return the name of it.
  27. nodeattrfunc: Function to decorate a node with attributes.
  28. The function shall accept one `node` object as
  29. argument and return the attributes.
  30. edgeattrfunc: Function to decorate a edge with attributes.
  31. The function shall accept two `node` objects as
  32. argument. The first the node and the second the child
  33. and return the attributes.
  34. edgetypefunc: Function to which gives the edge type.
  35. The function shall accept two `node` objects as
  36. argument. The first the node and the second the child
  37. and return the edge (i.e. '->').
  38. maxlevel (int): Limit export to this number of levels.
  39. >>> from anytree import Node
  40. >>> root = Node("root")
  41. >>> s0 = Node("sub0", parent=root, edge=2)
  42. >>> s0b = Node("sub0B", parent=s0, foo=4, edge=109)
  43. >>> s0a = Node("sub0A", parent=s0, edge="")
  44. >>> s1 = Node("sub1", parent=root, edge="")
  45. >>> s1a = Node("sub1A", parent=s1, edge=7)
  46. >>> s1b = Node("sub1B", parent=s1, edge=8)
  47. >>> s1c = Node("sub1C", parent=s1, edge=22)
  48. >>> s1ca = Node("sub1Ca", parent=s1c, edge=42)
  49. .. note:: If the node names are not unqiue, see :any:`UniqueDotExporter`.
  50. A directed graph:
  51. >>> from anytree.exporter import DotExporter
  52. >>> for line in DotExporter(root):
  53. ... print(line)
  54. digraph tree {
  55. "root";
  56. "sub0";
  57. "sub0B";
  58. "sub0A";
  59. "sub1";
  60. "sub1A";
  61. "sub1B";
  62. "sub1C";
  63. "sub1Ca";
  64. "root" -> "sub0";
  65. "root" -> "sub1";
  66. "sub0" -> "sub0B";
  67. "sub0" -> "sub0A";
  68. "sub1" -> "sub1A";
  69. "sub1" -> "sub1B";
  70. "sub1" -> "sub1C";
  71. "sub1C" -> "sub1Ca";
  72. }
  73. The resulting graph:
  74. .. image:: ../static/dotexporter0.png
  75. An undirected graph:
  76. >>> def nodenamefunc(node):
  77. ... return '%s:%s' % (node.name, node.depth)
  78. >>> def edgeattrfunc(node, child):
  79. ... return 'label="%s:%s"' % (node.name, child.name)
  80. >>> def edgetypefunc(node, child):
  81. ... return '--'
  82. >>> from anytree.exporter import DotExporter
  83. >>> for line in DotExporter(root, graph="graph",
  84. ... nodenamefunc=nodenamefunc,
  85. ... nodeattrfunc=lambda node: "shape=box",
  86. ... edgeattrfunc=edgeattrfunc,
  87. ... edgetypefunc=edgetypefunc):
  88. ... print(line)
  89. graph tree {
  90. "root:0" [shape=box];
  91. "sub0:1" [shape=box];
  92. "sub0B:2" [shape=box];
  93. "sub0A:2" [shape=box];
  94. "sub1:1" [shape=box];
  95. "sub1A:2" [shape=box];
  96. "sub1B:2" [shape=box];
  97. "sub1C:2" [shape=box];
  98. "sub1Ca:3" [shape=box];
  99. "root:0" -- "sub0:1" [label="root:sub0"];
  100. "root:0" -- "sub1:1" [label="root:sub1"];
  101. "sub0:1" -- "sub0B:2" [label="sub0:sub0B"];
  102. "sub0:1" -- "sub0A:2" [label="sub0:sub0A"];
  103. "sub1:1" -- "sub1A:2" [label="sub1:sub1A"];
  104. "sub1:1" -- "sub1B:2" [label="sub1:sub1B"];
  105. "sub1:1" -- "sub1C:2" [label="sub1:sub1C"];
  106. "sub1C:2" -- "sub1Ca:3" [label="sub1C:sub1Ca"];
  107. }
  108. The resulting graph:
  109. .. image:: ../static/dotexporter1.png
  110. To export custom node implementations or :any:`AnyNode`, please provide a proper `nodenamefunc`:
  111. >>> from anytree import AnyNode
  112. >>> root = AnyNode(id="root")
  113. >>> s0 = AnyNode(id="sub0", parent=root)
  114. >>> s0b = AnyNode(id="s0b", parent=s0)
  115. >>> s0a = AnyNode(id="s0a", parent=s0)
  116. >>> from anytree.exporter import DotExporter
  117. >>> for line in DotExporter(root, nodenamefunc=lambda n: n.id):
  118. ... print(line)
  119. digraph tree {
  120. "root";
  121. "sub0";
  122. "s0b";
  123. "s0a";
  124. "root" -> "sub0";
  125. "sub0" -> "s0b";
  126. "sub0" -> "s0a";
  127. }
  128. """
  129. self.node = node
  130. self.graph = graph
  131. self.name = name
  132. self.options = options
  133. self.indent = indent
  134. self.nodenamefunc = nodenamefunc
  135. self.nodeattrfunc = nodeattrfunc
  136. self.edgeattrfunc = edgeattrfunc
  137. self.edgetypefunc = edgetypefunc
  138. self.maxlevel = maxlevel
  139. def __iter__(self):
  140. # prepare
  141. indent = " " * self.indent
  142. nodenamefunc = self.nodenamefunc or self._default_nodenamefunc
  143. nodeattrfunc = self.nodeattrfunc or self._default_nodeattrfunc
  144. edgeattrfunc = self.edgeattrfunc or self._default_edgeattrfunc
  145. edgetypefunc = self.edgetypefunc or self._default_edgetypefunc
  146. return self.__iter(indent, nodenamefunc, nodeattrfunc, edgeattrfunc,
  147. edgetypefunc)
  148. @staticmethod
  149. def _default_nodenamefunc(node):
  150. return node.name
  151. @staticmethod
  152. def _default_nodeattrfunc(node):
  153. return None
  154. @staticmethod
  155. def _default_edgeattrfunc(node, child):
  156. return None
  157. @staticmethod
  158. def _default_edgetypefunc(node, child):
  159. return "->"
  160. def __iter(self, indent, nodenamefunc, nodeattrfunc, edgeattrfunc, edgetypefunc):
  161. yield "{self.graph} {self.name} {{".format(self=self)
  162. for option in self.__iter_options(indent):
  163. yield option
  164. for node in self.__iter_nodes(indent, nodenamefunc, nodeattrfunc):
  165. yield node
  166. for edge in self.__iter_edges(indent, nodenamefunc, edgeattrfunc, edgetypefunc):
  167. yield edge
  168. yield "}"
  169. def __iter_options(self, indent):
  170. options = self.options
  171. if options:
  172. for option in options:
  173. yield "%s%s" % (indent, option)
  174. def __iter_nodes(self, indent, nodenamefunc, nodeattrfunc):
  175. for node in PreOrderIter(self.node, maxlevel=self.maxlevel):
  176. nodename = nodenamefunc(node)
  177. nodeattr = nodeattrfunc(node)
  178. nodeattr = " [%s]" % nodeattr if nodeattr is not None else ""
  179. yield '%s"%s"%s;' % (indent, DotExporter.esc(nodename), nodeattr)
  180. def __iter_edges(self, indent, nodenamefunc, edgeattrfunc, edgetypefunc):
  181. maxlevel = self.maxlevel - 1 if self.maxlevel else None
  182. for node in PreOrderIter(self.node, maxlevel=maxlevel):
  183. nodename = nodenamefunc(node)
  184. for child in node.children:
  185. childname = nodenamefunc(child)
  186. edgeattr = edgeattrfunc(node, child)
  187. edgetype = edgetypefunc(node, child)
  188. edgeattr = " [%s]" % edgeattr if edgeattr is not None else ""
  189. yield '%s"%s" %s "%s"%s;' % (indent, DotExporter.esc(nodename), edgetype,
  190. DotExporter.esc(childname), edgeattr)
  191. def to_dotfile(self, filename):
  192. """
  193. Write graph to `filename`.
  194. >>> from anytree import Node
  195. >>> root = Node("root")
  196. >>> s0 = Node("sub0", parent=root)
  197. >>> s0b = Node("sub0B", parent=s0)
  198. >>> s0a = Node("sub0A", parent=s0)
  199. >>> s1 = Node("sub1", parent=root)
  200. >>> s1a = Node("sub1A", parent=s1)
  201. >>> s1b = Node("sub1B", parent=s1)
  202. >>> s1c = Node("sub1C", parent=s1)
  203. >>> s1ca = Node("sub1Ca", parent=s1c)
  204. >>> from anytree.exporter import DotExporter
  205. >>> DotExporter(root).to_dotfile("tree.dot")
  206. The generated file should be handed over to the `dot` tool from the
  207. http://www.graphviz.org/ package::
  208. $ dot tree.dot -T png -o tree.png
  209. """
  210. with codecs.open(filename, "w", "utf-8") as file:
  211. for line in self:
  212. file.write("%s\n" % line)
  213. def to_picture(self, filename):
  214. """
  215. Write graph to a temporary file and invoke `dot`.
  216. The output file type is automatically detected from the file suffix.
  217. *`graphviz` needs to be installed, before usage of this method.*
  218. """
  219. fileformat = path.splitext(filename)[1][1:]
  220. with NamedTemporaryFile("wb", delete=False) as dotfile:
  221. dotfilename = dotfile.name
  222. for line in self:
  223. dotfile.write(("%s\n" % line).encode("utf-8"))
  224. dotfile.flush()
  225. cmd = ["dot", dotfilename, "-T", fileformat, "-o", filename]
  226. check_call(cmd)
  227. try:
  228. remove(dotfilename)
  229. except Exception: # pragma: no cover
  230. msg = 'Could not remove temporary file %s' % dotfilename
  231. logging.getLogger(__name__).warn(msg)
  232. @staticmethod
  233. def esc(value):
  234. """Escape Strings."""
  235. return _RE_ESC.sub(lambda m: r"\%s" % m.group(0), six.text_type(value))
  236. class UniqueDotExporter(DotExporter):
  237. def __init__(self, node, graph="digraph", name="tree", options=None,
  238. indent=4, nodenamefunc=None, nodeattrfunc=None,
  239. edgeattrfunc=None, edgetypefunc=None):
  240. """
  241. Unqiue Dot Language Exporter.
  242. Handle trees with random or conflicting node names gracefully.
  243. Args:
  244. node (Node): start node.
  245. Keyword Args:
  246. graph: DOT graph type.
  247. name: DOT graph name.
  248. options: list of options added to the graph.
  249. indent (int): number of spaces for indent.
  250. nodenamefunc: Function to extract node name from `node` object.
  251. The function shall accept one `node` object as
  252. argument and return the name of it.
  253. nodeattrfunc: Function to decorate a node with attributes.
  254. The function shall accept one `node` object as
  255. argument and return the attributes.
  256. edgeattrfunc: Function to decorate a edge with attributes.
  257. The function shall accept two `node` objects as
  258. argument. The first the node and the second the child
  259. and return the attributes.
  260. edgetypefunc: Function to which gives the edge type.
  261. The function shall accept two `node` objects as
  262. argument. The first the node and the second the child
  263. and return the edge (i.e. '->').
  264. >>> from anytree import Node
  265. >>> root = Node("root")
  266. >>> s0 = Node("sub0", parent=root)
  267. >>> s0b = Node("s0", parent=s0)
  268. >>> s0a = Node("s0", parent=s0)
  269. >>> s1 = Node("sub1", parent=root)
  270. >>> s1a = Node("s1", parent=s1)
  271. >>> s1b = Node("s1", parent=s1)
  272. >>> s1c = Node("s1", parent=s1)
  273. >>> s1ca = Node("sub1Ca", parent=s1c)
  274. >>> from anytree.exporter import UniqueDotExporter
  275. >>> for line in UniqueDotExporter(root): # doctest: +SKIP
  276. ... print(line)
  277. digraph tree {
  278. "0x7f1bf2c9c510" [label="root"];
  279. "0x7f1bf2c9c5a0" [label="sub0"];
  280. "0x7f1bf2c9c630" [label="s0"];
  281. "0x7f1bf2c9c6c0" [label="s0"];
  282. "0x7f1bf2c9c750" [label="sub1"];
  283. "0x7f1bf2c9c7e0" [label="s1"];
  284. "0x7f1bf2c9c870" [label="s1"];
  285. "0x7f1bf2c9c900" [label="s1"];
  286. "0x7f1bf2c9c990" [label="sub1Ca"];
  287. "0x7f1bf2c9c510" -> "0x7f1bf2c9c5a0";
  288. "0x7f1bf2c9c510" -> "0x7f1bf2c9c750";
  289. "0x7f1bf2c9c5a0" -> "0x7f1bf2c9c630";
  290. "0x7f1bf2c9c5a0" -> "0x7f1bf2c9c6c0";
  291. "0x7f1bf2c9c750" -> "0x7f1bf2c9c7e0";
  292. "0x7f1bf2c9c750" -> "0x7f1bf2c9c870";
  293. "0x7f1bf2c9c750" -> "0x7f1bf2c9c900";
  294. "0x7f1bf2c9c900" -> "0x7f1bf2c9c990";
  295. }
  296. The resulting graph:
  297. .. image:: ../static/uniquedotexporter2.png
  298. To export custom node implementations or :any:`AnyNode`, please provide a proper `nodeattrfunc`:
  299. >>> from anytree import AnyNode
  300. >>> root = AnyNode(id="root")
  301. >>> s0 = AnyNode(id="sub0", parent=root)
  302. >>> s0b = AnyNode(id="s0", parent=s0)
  303. >>> s0a = AnyNode(id="s0", parent=s0)
  304. >>> from anytree.exporter import UniqueDotExporter
  305. >>> for line in UniqueDotExporter(root, nodeattrfunc=lambda n: 'label="%s"' % (n.id)): # doctest: +SKIP
  306. ... print(line)
  307. digraph tree {
  308. "0x7f5c70449af8" [label="root"];
  309. "0x7f5c70449bd0" [label="sub0"];
  310. "0x7f5c70449c60" [label="s0"];
  311. "0x7f5c70449cf0" [label="s0"];
  312. "0x7f5c70449af8" -> "0x7f5c70449bd0";
  313. "0x7f5c70449bd0" -> "0x7f5c70449c60";
  314. "0x7f5c70449bd0" -> "0x7f5c70449cf0";
  315. }
  316. """
  317. super(UniqueDotExporter, self).__init__(node, graph=graph, name=name, options=options, indent=indent,
  318. nodenamefunc=nodenamefunc, nodeattrfunc=nodeattrfunc,
  319. edgeattrfunc=edgeattrfunc, edgetypefunc=edgetypefunc)
  320. @staticmethod
  321. def _default_nodenamefunc(node):
  322. return hex(id(node))
  323. @staticmethod
  324. def _default_nodeattrfunc(node):
  325. return 'label="%s"' % (node.name)