search.py 7.2 KB

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