You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

305 lines
8.9 KiB

4 years ago
  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) and len(cont) == len(end)), (
  27. "'%s', '%s' and '%s' need to have equal length" % (vertical, cont,
  28. end))
  29. @property
  30. def empty(self):
  31. """Empty string as placeholder."""
  32. return ' ' * len(self.end)
  33. def __repr__(self):
  34. classname = self.__class__.__name__
  35. return "%s()" % classname
  36. class AsciiStyle(AbstractStyle):
  37. def __init__(self):
  38. """
  39. Ascii style.
  40. >>> from anytree import Node, RenderTree
  41. >>> root = Node("root")
  42. >>> s0 = Node("sub0", parent=root)
  43. >>> s0b = Node("sub0B", parent=s0)
  44. >>> s0a = Node("sub0A", parent=s0)
  45. >>> s1 = Node("sub1", parent=root)
  46. >>> print(RenderTree(root, style=AsciiStyle()))
  47. Node('/root')
  48. |-- Node('/root/sub0')
  49. | |-- Node('/root/sub0/sub0B')
  50. | +-- Node('/root/sub0/sub0A')
  51. +-- Node('/root/sub1')
  52. """
  53. super(AsciiStyle, self).__init__(u'| ', u'|-- ', u'+-- ')
  54. class ContStyle(AbstractStyle):
  55. def __init__(self):
  56. u"""
  57. Continued style, without gaps.
  58. >>> from anytree import Node, RenderTree
  59. >>> root = Node("root")
  60. >>> s0 = Node("sub0", parent=root)
  61. >>> s0b = Node("sub0B", parent=s0)
  62. >>> s0a = Node("sub0A", parent=s0)
  63. >>> s1 = Node("sub1", parent=root)
  64. >>> print(RenderTree(root, style=ContStyle()))
  65. Node('/root')
  66. Node('/root/sub0')
  67. Node('/root/sub0/sub0B')
  68. Node('/root/sub0/sub0A')
  69. Node('/root/sub1')
  70. """
  71. super(ContStyle, self).__init__(u'\u2502 ',
  72. u'\u251c\u2500\u2500 ',
  73. u'\u2514\u2500\u2500 ')
  74. class ContRoundStyle(AbstractStyle):
  75. def __init__(self):
  76. u"""
  77. Continued style, without gaps, round edges.
  78. >>> from anytree import Node, RenderTree
  79. >>> root = Node("root")
  80. >>> s0 = Node("sub0", parent=root)
  81. >>> s0b = Node("sub0B", parent=s0)
  82. >>> s0a = Node("sub0A", parent=s0)
  83. >>> s1 = Node("sub1", parent=root)
  84. >>> print(RenderTree(root, style=ContRoundStyle()))
  85. Node('/root')
  86. Node('/root/sub0')
  87. Node('/root/sub0/sub0B')
  88. Node('/root/sub0/sub0A')
  89. Node('/root/sub1')
  90. """
  91. super(ContRoundStyle, self).__init__(u'\u2502 ',
  92. u'\u251c\u2500\u2500 ',
  93. u'\u2570\u2500\u2500 ')
  94. class DoubleStyle(AbstractStyle):
  95. def __init__(self):
  96. u"""
  97. Double line style, without gaps.
  98. >>> from anytree import Node, RenderTree
  99. >>> root = Node("root")
  100. >>> s0 = Node("sub0", parent=root)
  101. >>> s0b = Node("sub0B", parent=s0)
  102. >>> s0a = Node("sub0A", parent=s0)
  103. >>> s1 = Node("sub1", parent=root)
  104. >>> print(RenderTree(root, style=DoubleStyle))
  105. Node('/root')
  106. Node('/root/sub0')
  107. Node('/root/sub0/sub0B')
  108. Node('/root/sub0/sub0A')
  109. Node('/root/sub1')
  110. """
  111. super(DoubleStyle, self).__init__(u'\u2551 ',
  112. u'\u2560\u2550\u2550 ',
  113. u'\u255a\u2550\u2550 ')
  114. @six.python_2_unicode_compatible
  115. class RenderTree(object):
  116. def __init__(self, node, style=ContStyle(), childiter=list):
  117. u"""
  118. Render tree starting at `node`.
  119. Keyword Args:
  120. style (AbstractStyle): Render Style.
  121. childiter: Child iterator.
  122. :any:`RenderTree` is an iterator, returning a tuple with 3 items:
  123. `pre`
  124. tree prefix.
  125. `fill`
  126. filling for multiline entries.
  127. `node`
  128. :any:`NodeMixin` object.
  129. It is up to the user to assemble these parts to a whole.
  130. >>> from anytree import Node, RenderTree
  131. >>> root = Node("root", lines=["c0fe", "c0de"])
  132. >>> s0 = Node("sub0", parent=root, lines=["ha", "ba"])
  133. >>> s0b = Node("sub0B", parent=s0, lines=["1", "2", "3"])
  134. >>> s0a = Node("sub0A", parent=s0, lines=["a", "b"])
  135. >>> s1 = Node("sub1", parent=root, lines=["Z"])
  136. Simple one line:
  137. >>> for pre, _, node in RenderTree(root):
  138. ... print("%s%s" % (pre, node.name))
  139. root
  140. sub0
  141. sub0B
  142. sub0A
  143. sub1
  144. Multiline:
  145. >>> for pre, fill, node in RenderTree(root):
  146. ... print("%s%s" % (pre, node.lines[0]))
  147. ... for line in node.lines[1:]:
  148. ... print("%s%s" % (fill, line))
  149. c0fe
  150. c0de
  151. ha
  152. ba
  153. 1
  154. 2
  155. 3
  156. a
  157. b
  158. Z
  159. The `childiter` is responsible for iterating over child nodes at the
  160. same level. An reversed order can be achived by using `reversed`.
  161. >>> for row in RenderTree(root, childiter=reversed):
  162. ... print("%s%s" % (row.pre, row.node.name))
  163. root
  164. sub1
  165. sub0
  166. sub0A
  167. sub0B
  168. Or writing your own sort function:
  169. >>> def mysort(items):
  170. ... return sorted(items, key=lambda item: item.name)
  171. >>> for row in RenderTree(root, childiter=mysort):
  172. ... print("%s%s" % (row.pre, row.node.name))
  173. root
  174. sub0
  175. sub0A
  176. sub0B
  177. sub1
  178. :any:`by_attr` simplifies attribute rendering and supports multiline:
  179. >>> print(RenderTree(root).by_attr())
  180. root
  181. sub0
  182. sub0B
  183. sub0A
  184. sub1
  185. >>> print(RenderTree(root).by_attr("lines"))
  186. c0fe
  187. c0de
  188. ha
  189. ba
  190. 1
  191. 2
  192. 3
  193. a
  194. b
  195. Z
  196. """
  197. if not isinstance(style, AbstractStyle):
  198. style = style()
  199. self.node = node
  200. self.style = style
  201. self.childiter = childiter
  202. def __iter__(self):
  203. return self.__next(self.node, tuple())
  204. def __next(self, node, continues):
  205. yield RenderTree.__item(node, continues, self.style)
  206. children = node.children
  207. if children:
  208. lastidx = len(children) - 1
  209. for idx, child in enumerate(self.childiter(children)):
  210. for grandchild in self.__next(child, continues + (idx != lastidx, )):
  211. yield grandchild
  212. @staticmethod
  213. def __item(node, continues, style):
  214. if not continues:
  215. return Row(u'', u'', node)
  216. else:
  217. items = [style.vertical if cont else style.empty for cont in continues]
  218. indent = ''.join(items[:-1])
  219. branch = style.cont if continues[-1] else style.end
  220. pre = indent + branch
  221. fill = ''.join(items)
  222. return Row(pre, fill, node)
  223. def __str__(self):
  224. lines = ["%s%r" % (pre, node) for pre, _, node in self]
  225. return "\n".join(lines)
  226. def __repr__(self):
  227. classname = self.__class__.__name__
  228. args = [repr(self.node),
  229. "style=%s" % repr(self.style),
  230. "childiter=%s" % repr(self.childiter)]
  231. return "%s(%s)" % (classname, ", ".join(args))
  232. def by_attr(self, attrname="name"):
  233. """Return rendered tree with node attribute `attrname`."""
  234. def get():
  235. for pre, fill, node in self:
  236. attr = getattr(node, attrname, "")
  237. if isinstance(attr, (list, tuple)):
  238. lines = attr
  239. else:
  240. lines = str(attr).split("\n")
  241. yield u"%s%s" % (pre, lines[0])
  242. for line in lines[1:]:
  243. yield u"%s%s" % (fill, line)
  244. return "\n".join(get())