Add containment for trees and graphs

penman/graph.py

@@ -4,7 +4,7 @@
 Data structures for Penman graphs and triples.
 """
 
-from typing import (Union, Optional, Mapping, List, Dict, Set, NamedTuple)
+from typing import (Union, Optional, Mapping, List, Dict, Set, NamedTuple, Tuple)
 from collections import defaultdict
 import copy
 
@@ -173,6 +173,20 @@ def __isub__(self, other):
         else:
             return NotImplemented
 
+    def __contains__(self, item: Union[str, Triple, Tuple[str, str, str]]):
+        """
+        Containment checking of either a string or a triple (tuple). In the case of a
+        string, will return True if the given string is a terminal instance in the graph.
+        If a triple or tuple, will return True if the triple exists in self.triples.
+        """
+        if isinstance(item, str):
+            terminals = [t[2] for t in self.instances()]
+            return item in terminals
+        elif isinstance(item, tuple):
+            return item in self.triples
+        else:
+            raise NotImplemented
+
     @property
     def top(self) -> Union[Variable, None]:
         """

penman/tree.py

@@ -40,6 +40,16 @@ def __str__(self) -> str:
         s = _format(self.node, 2)
         return f'Tree(\n  {s})'
 
+    def __contains__(self, item: str):
+        """
+        Test whether a given string is part of the tree as a terminal instance node.
+        """
+        if isinstance(item, str):
+            terminals = [target for _, (role, target) in self.walk() if role == "/" and is_atomic(target)]
+            return item in terminals
+        else:
+            raise NotImplemented
+
     def nodes(self) -> List[Node]:
         """
         Return the nodes in the tree as a flat list.

tests/test_graph.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import penman
+from penman import Triple
 
 Graph = penman.Graph
 
@@ -102,6 +103,17 @@ def test__isub__(self):
         assert g.top == 'b'
         assert g is original
 
+    def test__contains__(self):
+        g = Graph([('a', ':instance', 'alpha'),
+                   ('a', ':ARG', 'b'),
+                   ('b', ':instance', 'beta')])
+
+        assert "alpha" in g
+        assert "beta" in g
+        assert "b" not in g
+        assert ('a', ':ARG', 'b') in g
+        assert Triple('b', ':instance', 'beta') in g
+
     def test_top(self, x1):
         assert Graph([('a', ':instance', None)]).top == 'a'
         assert Graph(

tests/test_tree.py

@@ -1,4 +1,3 @@
-
 import pytest
 
 from penman import tree
@@ -47,6 +46,14 @@ def test__init__(self, empty_node, simple_node):
         assert t.node == simple_node
         assert t.metadata == {'snt': 'Alpha.'}
 
+    def test__contains__(self):
+        t = tree.Tree(('a', [('/', 'alpha'), (':ARG', ('b', [('/', 'beta')]))]))
+
+        assert "alpha" in t
+        assert "beta" in t
+        assert "b" not in t
+        assert ":ARG" not in t
+
     def test_nodes(self, one_arg_node, reentrant):
         t = tree.Tree(one_arg_node)
         assert t.nodes() == [one_arg_node, ('b', [('/', 'beta')])]
@@ -86,7 +93,6 @@ def test_walk(self, one_arg_node, reentrant):
         ]
 
     def test_reset_variables(self, one_arg_node, reentrant, var_instance):
-
         def _vars(t):
             return [v for v, _ in t.nodes()]