render.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368
  1. # -*- coding: utf-8 -*-
  2. """
  3. Tree Rendering.
  4. * :any:`RenderTree` using the following styles:
  5. * :any:`AsciiStyle`
  6. * :any:`ContStyle`
  7. * :any:`ContRoundStyle`
  8. * :any:`DoubleStyle`
  9. """
  10. import collections
  11. import six
  12. Row = collections.namedtuple("Row", ("pre", "fill", "node"))
  13. class AbstractStyle(object):
  14. def __init__(self, vertical, cont, end):
  15. """
  16. Tree Render Style.
  17. Args:
  18. vertical: Sign for vertical line.
  19. cont: Chars for a continued branch.
  20. end: Chars for the last branch.
  21. """
  22. super(AbstractStyle, self).__init__()
  23. self.vertical = vertical
  24. self.cont = cont
  25. self.end = end
  26. assert len(cont) == len(vertical) == len(end), "'%s', '%s' and '%s' need to have equal length" % (
  27. vertical,
  28. cont,
  29. end,
  30. )
  31. @property
  32. def empty(self):
  33. """Empty string as placeholder."""
  34. return ' ' * len(self.end)
  35. def __repr__(self):
  36. classname = self.__class__.__name__
  37. return "%s()" % classname
  38. class AsciiStyle(AbstractStyle):
  39. def __init__(self):
  40. """
  41. Ascii style.
  42. >>> from anytree import Node, RenderTree
  43. >>> root = Node("root")
  44. >>> s0 = Node("sub0", parent=root)
  45. >>> s0b = Node("sub0B", parent=s0)
  46. >>> s0a = Node("sub0A", parent=s0)
  47. >>> s1 = Node("sub1", parent=root)
  48. >>> print(RenderTree(root, style=AsciiStyle()))
  49. Node('/root')
  50. |-- Node('/root/sub0')
  51. | |-- Node('/root/sub0/sub0B')
  52. | +-- Node('/root/sub0/sub0A')
  53. +-- Node('/root/sub1')
  54. """
  55. super(AsciiStyle, self).__init__(u'| ', u'|-- ', u'+-- ')
  56. class ContStyle(AbstractStyle):
  57. def __init__(self):
  58. u"""
  59. Continued style, without gaps.
  60. >>> from anytree import Node, RenderTree
  61. >>> root = Node("root")
  62. >>> s0 = Node("sub0", parent=root)
  63. >>> s0b = Node("sub0B", parent=s0)
  64. >>> s0a = Node("sub0A", parent=s0)
  65. >>> s1 = Node("sub1", parent=root)
  66. >>> print(RenderTree(root, style=ContStyle()))
  67. Node('/root')
  68. ├── Node('/root/sub0')
  69. │ ├── Node('/root/sub0/sub0B')
  70. │ └── Node('/root/sub0/sub0A')
  71. └── Node('/root/sub1')
  72. """
  73. super(ContStyle, self).__init__(u'\u2502 ',
  74. u'\u251c\u2500\u2500 ',
  75. u'\u2514\u2500\u2500 ')
  76. class ContRoundStyle(AbstractStyle):
  77. def __init__(self):
  78. u"""
  79. Continued style, without gaps, round edges.
  80. >>> from anytree import Node, RenderTree
  81. >>> root = Node("root")
  82. >>> s0 = Node("sub0", parent=root)
  83. >>> s0b = Node("sub0B", parent=s0)
  84. >>> s0a = Node("sub0A", parent=s0)
  85. >>> s1 = Node("sub1", parent=root)
  86. >>> print(RenderTree(root, style=ContRoundStyle()))
  87. Node('/root')
  88. ├── Node('/root/sub0')
  89. │ ├── Node('/root/sub0/sub0B')
  90. │ ╰── Node('/root/sub0/sub0A')
  91. ╰── Node('/root/sub1')
  92. """
  93. super(ContRoundStyle, self).__init__(u'\u2502 ',
  94. u'\u251c\u2500\u2500 ',
  95. u'\u2570\u2500\u2500 ')
  96. class DoubleStyle(AbstractStyle):
  97. def __init__(self):
  98. u"""
  99. Double line style, without gaps.
  100. >>> from anytree import Node, RenderTree
  101. >>> root = Node("root")
  102. >>> s0 = Node("sub0", parent=root)
  103. >>> s0b = Node("sub0B", parent=s0)
  104. >>> s0a = Node("sub0A", parent=s0)
  105. >>> s1 = Node("sub1", parent=root)
  106. >>> print(RenderTree(root, style=DoubleStyle))
  107. Node('/root')
  108. ╠══ Node('/root/sub0')
  109. ║ ╠══ Node('/root/sub0/sub0B')
  110. ║ ╚══ Node('/root/sub0/sub0A')
  111. ╚══ Node('/root/sub1')
  112. """
  113. super(DoubleStyle, self).__init__(u'\u2551 ',
  114. u'\u2560\u2550\u2550 ',
  115. u'\u255a\u2550\u2550 ')
  116. @six.python_2_unicode_compatible
  117. class RenderTree(object):
  118. def __init__(self, node, style=ContStyle(), childiter=list, maxlevel=None):
  119. u"""
  120. Render tree starting at `node`.
  121. Keyword Args:
  122. style (AbstractStyle): Render Style.
  123. childiter: Child iterator.
  124. maxlevel: Limit rendering to this depth.
  125. :any:`RenderTree` is an iterator, returning a tuple with 3 items:
  126. `pre`
  127. tree prefix.
  128. `fill`
  129. filling for multiline entries.
  130. `node`
  131. :any:`NodeMixin` object.
  132. It is up to the user to assemble these parts to a whole.
  133. >>> from anytree import Node, RenderTree
  134. >>> root = Node("root", lines=["c0fe", "c0de"])
  135. >>> s0 = Node("sub0", parent=root, lines=["ha", "ba"])
  136. >>> s0b = Node("sub0B", parent=s0, lines=["1", "2", "3"])
  137. >>> s0a = Node("sub0A", parent=s0, lines=["a", "b"])
  138. >>> s1 = Node("sub1", parent=root, lines=["Z"])
  139. Simple one line:
  140. >>> for pre, _, node in RenderTree(root):
  141. ... print("%s%s" % (pre, node.name))
  142. root
  143. ├── sub0
  144. │ ├── sub0B
  145. │ └── sub0A
  146. └── sub1
  147. Multiline:
  148. >>> for pre, fill, node in RenderTree(root):
  149. ... print("%s%s" % (pre, node.lines[0]))
  150. ... for line in node.lines[1:]:
  151. ... print("%s%s" % (fill, line))
  152. c0fe
  153. c0de
  154. ├── ha
  155. │ ba
  156. │ ├── 1
  157. │ │ 2
  158. │ │ 3
  159. │ └── a
  160. │ b
  161. └── Z
  162. `maxlevel` limits the depth of the tree:
  163. >>> print(RenderTree(root, maxlevel=2))
  164. Node('/root', lines=['c0fe', 'c0de'])
  165. ├── Node('/root/sub0', lines=['ha', 'ba'])
  166. └── Node('/root/sub1', lines=['Z'])
  167. The `childiter` is responsible for iterating over child nodes at the
  168. same level. An reversed order can be achived by using `reversed`.
  169. >>> for row in RenderTree(root, childiter=reversed):
  170. ... print("%s%s" % (row.pre, row.node.name))
  171. root
  172. ├── sub1
  173. └── sub0
  174. ├── sub0A
  175. └── sub0B
  176. Or writing your own sort function:
  177. >>> def mysort(items):
  178. ... return sorted(items, key=lambda item: item.name)
  179. >>> for row in RenderTree(root, childiter=mysort):
  180. ... print("%s%s" % (row.pre, row.node.name))
  181. root
  182. ├── sub0
  183. │ ├── sub0A
  184. │ └── sub0B
  185. └── sub1
  186. :any:`by_attr` simplifies attribute rendering and supports multiline:
  187. >>> print(RenderTree(root).by_attr())
  188. root
  189. ├── sub0
  190. │ ├── sub0B
  191. │ └── sub0A
  192. └── sub1
  193. >>> print(RenderTree(root).by_attr("lines"))
  194. c0fe
  195. c0de
  196. ├── ha
  197. │ ba
  198. │ ├── 1
  199. │ │ 2
  200. │ │ 3
  201. │ └── a
  202. │ b
  203. └── Z
  204. And can be a function:
  205. >>> print(RenderTree(root).by_attr(lambda n: " ".join(n.lines)))
  206. c0fe c0de
  207. ├── ha ba
  208. │ ├── 1 2 3
  209. │ └── a b
  210. └── Z
  211. """
  212. if not isinstance(style, AbstractStyle):
  213. style = style()
  214. self.node = node
  215. self.style = style
  216. self.childiter = childiter
  217. self.maxlevel = maxlevel
  218. def __iter__(self):
  219. return self.__next(self.node, tuple())
  220. def __next(self, node, continues, level=0):
  221. yield RenderTree.__item(node, continues, self.style)
  222. children = node.children
  223. level += 1
  224. if children and (self.maxlevel is None or level < self.maxlevel):
  225. children = self.childiter(children)
  226. for child, is_last in _is_last(children):
  227. for grandchild in self.__next(child, continues + (not is_last, ), level=level):
  228. yield grandchild
  229. @staticmethod
  230. def __item(node, continues, style):
  231. if not continues:
  232. return Row(u'', u'', node)
  233. else:
  234. items = [style.vertical if cont else style.empty for cont in continues]
  235. indent = ''.join(items[:-1])
  236. branch = style.cont if continues[-1] else style.end
  237. pre = indent + branch
  238. fill = ''.join(items)
  239. return Row(pre, fill, node)
  240. def __str__(self):
  241. lines = ["%s%r" % (pre, node) for pre, _, node in self]
  242. return "\n".join(lines)
  243. def __repr__(self):
  244. classname = self.__class__.__name__
  245. args = [repr(self.node),
  246. "style=%s" % repr(self.style),
  247. "childiter=%s" % repr(self.childiter)]
  248. return "%s(%s)" % (classname, ", ".join(args))
  249. def by_attr(self, attrname="name"):
  250. u"""
  251. Return rendered tree with node attribute `attrname`.
  252. >>> from anytree import AnyNode, RenderTree
  253. >>> root = AnyNode(id="root")
  254. >>> s0 = AnyNode(id="sub0", parent=root)
  255. >>> s0b = AnyNode(id="sub0B", parent=s0, foo=4, bar=109)
  256. >>> s0a = AnyNode(id="sub0A", parent=s0)
  257. >>> s1 = AnyNode(id="sub1", parent=root)
  258. >>> s1a = AnyNode(id="sub1A", parent=s1)
  259. >>> s1b = AnyNode(id="sub1B", parent=s1, bar=8)
  260. >>> s1c = AnyNode(id="sub1C", parent=s1)
  261. >>> s1ca = AnyNode(id="sub1Ca", parent=s1c)
  262. >>> print(RenderTree(root).by_attr('id'))
  263. root
  264. ├── sub0
  265. │ ├── sub0B
  266. │ └── sub0A
  267. └── sub1
  268. ├── sub1A
  269. ├── sub1B
  270. └── sub1C
  271. └── sub1Ca
  272. """
  273. def get():
  274. for pre, fill, node in self:
  275. attr = attrname(node) if callable(attrname) else getattr(node, attrname, "")
  276. if isinstance(attr, (list, tuple)):
  277. lines = attr
  278. else:
  279. lines = str(attr).split("\n")
  280. yield u"%s%s" % (pre, lines[0])
  281. for line in lines[1:]:
  282. yield u"%s%s" % (fill, line)
  283. return "\n".join(get())
  284. def _is_last(iterable):
  285. iter_ = iter(iterable)
  286. try:
  287. nextitem = next(iter_)
  288. except StopIteration:
  289. pass
  290. else:
  291. item = nextitem
  292. while True:
  293. try:
  294. nextitem = next(iter_)
  295. yield item, False
  296. except StopIteration:
  297. yield nextitem, True
  298. break
  299. item = nextitem