add __contains__ method

inducer · Jan 13, 2023 · 76cb176 · 76cb176
1 parent 4d60944
commit 76cb176
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py
@@ -125,6 +125,9 @@ def is_leaf(self, node: NodeT) -> bool:
  def is_a_node(self, node: NodeT) -> bool:
  return node in self._child_to_parent
 
+ def __contains__(self, node: NodeT) -> bool:
+ return self.is_a_node(node)
+
  def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]":
  """
  Returns a :class:`Tree` with added node *node* having a parent