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.

237 lines
7.0 KiB

4 years ago
  1. """Node Searching."""
  2. from anytree.iterators import PreOrderIter
  3. def findall(node, filter_=None, stop=None, maxlevel=None, mincount=None, maxcount=None):
  4. """
  5. Search nodes matching `filter_` but stop at `maxlevel` or `stop`.
  6. Return tuple with matching nodes.
  7. Args:
  8. node: top node, start searching.
  9. Keyword Args:
  10. filter_: function called with every `node` as argument, `node` is returned if `True`.
  11. stop: stop iteration at `node` if `stop` function returns `True` for `node`.
  12. maxlevel (int): maximum decending in the node hierarchy.
  13. mincount (int): minimum number of nodes.
  14. maxcount (int): maximum number of nodes.
  15. Example tree:
  16. >>> from anytree import Node, RenderTree, AsciiStyle
  17. >>> f = Node("f")
  18. >>> b = Node("b", parent=f)
  19. >>> a = Node("a", parent=b)
  20. >>> d = Node("d", parent=b)
  21. >>> c = Node("c", parent=d)
  22. >>> e = Node("e", parent=d)
  23. >>> g = Node("g", parent=f)
  24. >>> i = Node("i", parent=g)
  25. >>> h = Node("h", parent=i)
  26. >>> print(RenderTree(f, style=AsciiStyle()).by_attr())
  27. f
  28. |-- b
  29. | |-- a
  30. | +-- d
  31. | |-- c
  32. | +-- e
  33. +-- g
  34. +-- i
  35. +-- h
  36. >>> findall(f, filter_=lambda node: node.name in ("a", "b"))
  37. (Node('/f/b'), Node('/f/b/a'))
  38. >>> findall(f, filter_=lambda node: d in node.path)
  39. (Node('/f/b/d'), Node('/f/b/d/c'), Node('/f/b/d/e'))
  40. The number of matches can be limited:
  41. >>> findall(f, filter_=lambda node: d in node.path, mincount=4) # doctest: +ELLIPSIS
  42. Traceback (most recent call last):
  43. ...
  44. anytree.search.CountError: Expecting at least 4 elements, but found 3. ... Node('/f/b/d/e'))
  45. >>> findall(f, filter_=lambda node: d in node.path, maxcount=2) # doctest: +ELLIPSIS
  46. Traceback (most recent call last):
  47. ...
  48. anytree.search.CountError: Expecting 2 elements at maximum, but found 3. ... Node('/f/b/d/e'))
  49. """
  50. return _findall(node, filter_=filter_, stop=stop,
  51. maxlevel=maxlevel, mincount=mincount, maxcount=maxcount)
  52. def findall_by_attr(node, value, name="name", maxlevel=None, mincount=None, maxcount=None):
  53. """
  54. Search nodes with attribute `name` having `value` but stop at `maxlevel`.
  55. Return tuple with matching nodes.
  56. Args:
  57. node: top node, start searching.
  58. value: value which need to match
  59. Keyword Args:
  60. name (str): attribute name need to match
  61. maxlevel (int): maximum decending in the node hierarchy.
  62. mincount (int): minimum number of nodes.
  63. maxcount (int): maximum number of nodes.
  64. Example tree:
  65. >>> from anytree import Node, RenderTree, AsciiStyle
  66. >>> f = Node("f")
  67. >>> b = Node("b", parent=f)
  68. >>> a = Node("a", parent=b)
  69. >>> d = Node("d", parent=b)
  70. >>> c = Node("c", parent=d)
  71. >>> e = Node("e", parent=d)
  72. >>> g = Node("g", parent=f)
  73. >>> i = Node("i", parent=g)
  74. >>> h = Node("h", parent=i)
  75. >>> print(RenderTree(f, style=AsciiStyle()).by_attr())
  76. f
  77. |-- b
  78. | |-- a
  79. | +-- d
  80. | |-- c
  81. | +-- e
  82. +-- g
  83. +-- i
  84. +-- h
  85. >>> findall_by_attr(f, "d")
  86. (Node('/f/b/d'),)
  87. """
  88. return _findall(node, filter_=lambda n: _filter_by_name(n, name, value),
  89. maxlevel=maxlevel, mincount=mincount, maxcount=maxcount)
  90. def find(node, filter_=None, stop=None, maxlevel=None):
  91. """
  92. Search for *single* node matching `filter_` but stop at `maxlevel` or `stop`.
  93. Return matching node.
  94. Args:
  95. node: top node, start searching.
  96. Keyword Args:
  97. filter_: function called with every `node` as argument, `node` is returned if `True`.
  98. stop: stop iteration at `node` if `stop` function returns `True` for `node`.
  99. maxlevel (int): maximum decending in the node hierarchy.
  100. Example tree:
  101. >>> from anytree import Node, RenderTree, AsciiStyle
  102. >>> f = Node("f")
  103. >>> b = Node("b", parent=f)
  104. >>> a = Node("a", parent=b)
  105. >>> d = Node("d", parent=b)
  106. >>> c = Node("c", parent=d)
  107. >>> e = Node("e", parent=d)
  108. >>> g = Node("g", parent=f)
  109. >>> i = Node("i", parent=g)
  110. >>> h = Node("h", parent=i)
  111. >>> print(RenderTree(f, style=AsciiStyle()).by_attr())
  112. f
  113. |-- b
  114. | |-- a
  115. | +-- d
  116. | |-- c
  117. | +-- e
  118. +-- g
  119. +-- i
  120. +-- h
  121. >>> find(f, lambda node: node.name == "d")
  122. Node('/f/b/d')
  123. >>> find(f, lambda node: node.name == "z")
  124. >>> find(f, lambda node: b in node.path) # doctest: +ELLIPSIS
  125. Traceback (most recent call last):
  126. ...
  127. anytree.search.CountError: Expecting 1 elements at maximum, but found 5. (Node('/f/b')... Node('/f/b/d/e'))
  128. """
  129. return _find(node, filter_=filter_, stop=stop, maxlevel=maxlevel)
  130. def find_by_attr(node, value, name="name", maxlevel=None):
  131. """
  132. Search for *single* node with attribute `name` having `value` but stop at `maxlevel`.
  133. Return tuple with matching nodes.
  134. Args:
  135. node: top node, start searching.
  136. value: value which need to match
  137. Keyword Args:
  138. name (str): attribute name need to match
  139. maxlevel (int): maximum decending in the node hierarchy.
  140. Example tree:
  141. >>> from anytree import Node, RenderTree, AsciiStyle
  142. >>> f = Node("f")
  143. >>> b = Node("b", parent=f)
  144. >>> a = Node("a", parent=b)
  145. >>> d = Node("d", parent=b)
  146. >>> c = Node("c", parent=d, foo=4)
  147. >>> e = Node("e", parent=d)
  148. >>> g = Node("g", parent=f)
  149. >>> i = Node("i", parent=g)
  150. >>> h = Node("h", parent=i)
  151. >>> print(RenderTree(f, style=AsciiStyle()).by_attr())
  152. f
  153. |-- b
  154. | |-- a
  155. | +-- d
  156. | |-- c
  157. | +-- e
  158. +-- g
  159. +-- i
  160. +-- h
  161. >>> find_by_attr(f, "d")
  162. Node('/f/b/d')
  163. >>> find_by_attr(f, name="foo", value=4)
  164. Node('/f/b/d/c', foo=4)
  165. >>> find_by_attr(f, name="foo", value=8)
  166. """
  167. return _find(node, filter_=lambda n: _filter_by_name(n, name, value),
  168. maxlevel=maxlevel)
  169. def _find(node, filter_, stop=None, maxlevel=None):
  170. items = _findall(node, filter_, stop=stop, maxlevel=maxlevel, maxcount=1)
  171. return items[0] if items else None
  172. def _findall(node, filter_, stop=None, maxlevel=None, mincount=None, maxcount=None):
  173. result = tuple(PreOrderIter(node, filter_, stop, maxlevel))
  174. resultlen = len(result)
  175. if mincount is not None and resultlen < mincount:
  176. msg = "Expecting at least %d elements, but found %d."
  177. raise CountError(msg % (mincount, resultlen), result)
  178. if maxcount is not None and resultlen > maxcount:
  179. msg = "Expecting %d elements at maximum, but found %d."
  180. raise CountError(msg % (maxcount, resultlen), result)
  181. return result
  182. def _filter_by_name(node, name, value):
  183. try:
  184. return getattr(node, name) == value
  185. except AttributeError:
  186. return False
  187. class CountError(RuntimeError):
  188. def __init__(self, msg, result):
  189. """Error raised on `mincount` or `maxcount` mismatch."""
  190. if result:
  191. msg += " " + repr(result)
  192. super(CountError, self).__init__(msg)