Add more tests

Author: cthoyt

src/curies/triples/__init__.py

@@ -139,7 +139,8 @@
     exclude_subject_prefixes,
     exclude_triples,
     keep_object_prefixes,
-    keep_prefixes,
+    keep_prefixes_both,
+    keep_prefixes_either,
     keep_references_both,
     keep_references_either,
     keep_subject_prefixes,
@@ -162,7 +163,8 @@
     "exclude_triples",
     "hash_triple",
     "keep_object_prefixes",
-    "keep_prefixes",
+    "keep_prefixes_both",
+    "keep_prefixes_either",
     "keep_references_both",
     "keep_references_either",
     "keep_subject_prefixes",

src/curies/triples/filters.py

@@ -16,7 +16,8 @@
     "exclude_subject_prefixes",
     "exclude_triples",
     "keep_object_prefixes",
-    "keep_prefixes",
+    "keep_prefixes_both",
+    "keep_prefixes_either",
     "keep_references_both",
     "keep_references_either",
     "keep_subject_prefixes",
@@ -29,7 +30,7 @@
 def _filter(
     func: TriplePredicate[TripleType], triples: Iterable[TripleType], progress: bool = False
 ) -> Iterable[TripleType]:
-    if progress:
+    if progress:  # pragma: no cover
         try:
             from tqdm import tqdm
         except ImportError:
@@ -39,7 +40,7 @@ def _filter(
     return filter(func, triples)
 
 
-def keep_prefixes(
+def keep_prefixes_both(
     triples: Iterable[TripleType], prefixes: Iterable[str], *, progress: bool = False
 ) -> Iterable[TripleType]:
     """Keep triples whose subjects' and objects' prefixes are in the given prefixes.
@@ -57,12 +58,12 @@ def keep_prefixes(
     >>> m1 = Triple.from_curies(c1, exact_match.curie, c2)
     >>> m2 = Triple.from_curies(c2, exact_match.curie, c3)
     >>> m3 = Triple.from_curies(c1, exact_match.curie, c3)
-    >>> assert list(keep_prefixes([m1, m2, m3], {"DOID", "mesh"})) == [m1]
+    >>> assert list(keep_prefixes_both([m1, m2, m3], {"DOID", "mesh"})) == [m1]
     """
-    return _filter(_keep_prefixes_filter(prefixes), triples, progress=progress)
+    return _filter(_keep_prefixes_both_filter(prefixes), triples, progress=progress)
 
 
-def _keep_prefixes_filter(prefixes: Iterable[str]) -> TriplePredicate[TripleType]:
+def _keep_prefixes_both_filter(prefixes: Iterable[str]) -> TriplePredicate[TripleType]:
     prefixes = set(prefixes)
     if len(prefixes) < 2:
         raise ValueError
@@ -73,6 +74,44 @@ def _func(triple: TripleType) -> bool:
     return _func
 
 
+def keep_prefixes_either(
+    triples: Iterable[TripleType], prefixes: str | Iterable[str], *, progress: bool = False
+) -> Iterable[TripleType]:
+    """Keep triples whose subjects' and objects' prefixes are in the given prefixes.
+
+    :param triples: An iterable of triples
+    :param prefixes: A set of prefixes to use for filtering the triples
+    :param progress: Should a progress bar be shown?
+
+    :returns: A sub-iterable of triples whose subjects' and objects'
+        prefixes are in the given prefixes
+
+    >>> from curies import Triple
+    >>> from curies.vocabulary import exact_match
+    >>> c1, c2, c3 = "DOID:0050577", "mesh:C562966", "umls:C4551571"
+    >>> m1 = Triple.from_curies(c1, exact_match.curie, c2)
+    >>> m2 = Triple.from_curies(c2, exact_match.curie, c3)
+    >>> m3 = Triple.from_curies(c1, exact_match.curie, c3)
+    >>> assert list(keep_prefixes_either([m1, m2, m3], {"DOID", "mesh"})) == [m1]
+    """
+    return _filter(_keep_prefixes_either_filter(prefixes), triples, progress=progress)
+
+
+def _keep_prefixes_either_filter(prefixes: Iterable[str]) -> TriplePredicate[TripleType]:
+    if isinstance(prefixes, str):
+
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix == prefixes or triple.object.prefix == prefixes
+
+    else:
+        prefixes = set(prefixes)
+
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix in prefixes or triple.object.prefix in prefixes
+
+    return _func
+
+
 def keep_subject_prefixes(
     triples: Iterable[TripleType], prefixes: str | Iterable[str], *, progress: bool = False
 ) -> Iterable[TripleType]:

tests/test_triples/test_filter.py

@@ -12,14 +12,16 @@
     exclude_subject_prefixes,
     exclude_triples,
     keep_object_prefixes,
-    keep_prefixes,
+    keep_prefixes_both,
+    keep_prefixes_either,
     keep_references_both,
     keep_references_either,
     keep_subject_prefixes,
     keep_triples_by_hash,
 )
 from curies.vocabulary import exact_match, subclass_of
 
+p1, p2, p3 = "DOID", "mesh", "umls"
 c1, c2, c3, c4 = "DOID:0050577", "mesh:C562966", "umls:C4551571", "DOID:225"
 r1, r2, r3, r4 = (Reference.from_curie(c) for c in [c1, c2, c3, c4])
 m1 = Triple.from_curies(c1, exact_match.curie, c2)
@@ -51,18 +53,21 @@ def assert_triple_lists(self, expected: list[Triple], actual: Iterable[Triple])
 
     def test_exclude_object_prefixes(self) -> None:
         """Test excluding object prefixes."""
+        self.assert_triple_lists([m1], exclude_object_prefixes(M123, "umls"))
         self.assert_triple_lists([m1], exclude_object_prefixes(M123, {"umls"}))
         self.assert_triple_lists([m2, m3], exclude_object_prefixes(M123, {"mesh"}))
         self.assert_triple_lists(M123, exclude_object_prefixes(M123, {"DOID"}))
 
     def test_exclude_prefixes(self) -> None:
         """Test excluding prefixes."""
+        self.assert_triple_lists([m1], exclude_prefixes(M123, "umls"))
         self.assert_triple_lists([m1], exclude_prefixes(M123, {"umls"}))
         self.assert_triple_lists([m2], exclude_prefixes(M123, {"DOID"}))
         self.assert_triple_lists([m3], exclude_prefixes(M123, {"mesh"}))
 
     def test_exclude_subject_prefixes(self) -> None:
         """Test excluding subject prefixes."""
+        self.assert_triple_lists([m2], exclude_subject_prefixes(M123, "DOID"))
         self.assert_triple_lists([m2], exclude_subject_prefixes(M123, {"DOID"}))
         self.assert_triple_lists(M123, exclude_subject_prefixes(M123, {"umls"}))
         self.assert_triple_lists([m1, m3], exclude_subject_prefixes(M123, {"mesh"}))
@@ -78,18 +83,33 @@ def test_exclude_triples(self) -> None:
 
     def test_keep_object_prefixes(self) -> None:
         """Test keeping object prefixes."""
+        self.assert_triple_lists([m2, m3], keep_object_prefixes(M123, "umls"))
         self.assert_triple_lists([m2, m3], keep_object_prefixes(M123, {"umls"}))
 
-    def test_keep_prefixes(self) -> None:
+    def test_keep_prefixes_both(self) -> None:
         """Test keeping prefixes."""
-        self.assert_triple_lists([], keep_prefixes(M123, {"NOPE", "also nope"}))
-        self.assert_triple_lists([m1], keep_prefixes(M123, {"DOID", "mesh"}))
-        self.assert_triple_lists([m1], keep_prefixes(M123, {"DOID", "umls"}))
-        self.assert_triple_lists([m1], keep_prefixes(M123, {"mesh", "umls"}))
-        self.assert_triple_lists(M123, keep_prefixes(M123, {"DOID", "umls", "mesh"}))
+        with self.assertRaises(ValueError):
+            keep_prefixes_both(M123, {"DOID"})
+        self.assert_triple_lists([], keep_prefixes_both(M123, {"NOPE", "also nope"}))
+        self.assert_triple_lists([m1], keep_prefixes_both(M123, {"DOID", "mesh"}))
+        self.assert_triple_lists([m3], keep_prefixes_both(M123, {"DOID", "umls"}))
+        self.assert_triple_lists([m2], keep_prefixes_both(M123, {"mesh", "umls"}))
+        self.assert_triple_lists(M123, keep_prefixes_both(M123, {"DOID", "umls", "mesh"}))
+
+    def test_keep_prefixes_either(self) -> None:
+        """Test keeping prefixes."""
+        self.assert_triple_lists([], keep_prefixes_either(M123, {"NOPE", "also nope"}))
+        self.assert_triple_lists([m1, m2], keep_prefixes_either(M123, "mesh"))
+        self.assert_triple_lists([m1, m3], keep_prefixes_either(M123, "DOID"))
+        self.assert_triple_lists([m2, m3], keep_prefixes_either(M123, "umls"))
+        self.assert_triple_lists([m1, m2, m3], keep_prefixes_either(M123, {"DOID", "mesh"}))
+        self.assert_triple_lists([m1, m2, m3], keep_prefixes_either(M123, {"DOID", "umls"}))
+        self.assert_triple_lists([m1, m2, m3], keep_prefixes_either(M123, {"mesh", "umls"}))
+        self.assert_triple_lists(M123, keep_prefixes_either(M123, {"DOID", "umls", "mesh"}))
 
     def test_keep_subject_prefixes(self) -> None:
         """Test keeping subject prefixes."""
+        self.assert_triple_lists([m1, m3], keep_subject_prefixes(M123, "DOID"))
         self.assert_triple_lists([m1, m3], keep_subject_prefixes(M123, {"DOID"}))
 
     def test_keep_triple_by_hash(self) -> None:
@@ -106,6 +126,7 @@ def test_keep_triple_by_hash(self) -> None:
 
     def test_keep_references_either(self) -> None:
         """Test keeping references."""
+        self.assert_triple_lists([m1, m3], keep_references_either(M123, r1))
         self.assert_triple_lists([m1, m3], keep_references_either(M123, [r1]))
         self.assert_triple_lists([m1, m2], keep_references_either(M123, [r2]))
         self.assert_triple_lists([m2, m3], keep_references_either(M123, [r3]))
@@ -115,13 +136,16 @@ def test_keep_references_either(self) -> None:
 
     def test_keep_references_both(self) -> None:
         """Test keeping references."""
+        with self.assertRaises(ValueError):
+            keep_references_both(M123, [r1])
         self.assert_triple_lists([m1], keep_references_both(M123, [r1, r2]))
         self.assert_triple_lists([m2], keep_references_both(M123, [r2, r3]))
         self.assert_triple_lists([m3], keep_references_both(M123, [r1, r3]))
         self.assert_triple_lists(M123, keep_references_both(M123, [r1, r2, r3]))
 
     def test_exclude_references(self) -> None:
         """Test exclude references."""
+        self.assert_triple_lists([m2], exclude_references(M123, r1))
         self.assert_triple_lists([m2], exclude_references(M123, [r1]))
         self.assert_triple_lists([m3], exclude_references(M123, [r2]))
         self.assert_triple_lists([m1], exclude_references(M123, [r3]))