Improve typing and testing

Author: cthoyt

src/curies/triples/filters.py

@@ -26,12 +26,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _cleanup_prefixes(prefixes: str | Iterable[str]) -> set[str]:
-    if isinstance(prefixes, str):
-        return {prefixes}
-    return set(prefixes)
-
-
 def _filter(
     func: TriplePredicate[TripleType], triples: Iterable[TripleType], progress: bool = False
 ) -> Iterable[TripleType]:
@@ -46,7 +40,7 @@ def _filter(
 
 
 def keep_prefixes(
-    triples: Iterable[TripleType], prefixes: str | Iterable[str], *, progress: bool = False
+    triples: Iterable[TripleType], prefixes: Iterable[str], *, progress: bool = False
 ) -> Iterable[TripleType]:
     """Keep triples whose subjects' and objects' prefixes are in the given prefixes.
 
@@ -68,8 +62,10 @@ def keep_prefixes(
     return _filter(_keep_prefixes_filter(prefixes), triples, progress=progress)
 
 
-def _keep_prefixes_filter(prefixes: str | Iterable[str]) -> TriplePredicate[TripleType]:
-    prefixes = _cleanup_prefixes(prefixes)
+def _keep_prefixes_filter(prefixes: Iterable[str]) -> TriplePredicate[TripleType]:
+    prefixes = set(prefixes)
+    if len(prefixes) < 2:
+        raise ValueError
 
     def _func(triple: TripleType) -> bool:
         return triple.subject.prefix in prefixes and triple.object.prefix in prefixes
@@ -101,10 +97,16 @@ def keep_subject_prefixes(
 
 
 def _keep_subject_prefixes_filter(prefixes: str | Iterable[str]) -> TriplePredicate[TripleType]:
-    prefixes = _cleanup_prefixes(prefixes)
+    if isinstance(prefixes, str):
 
-    def _func(triple: TripleType) -> bool:
-        return triple.subject.prefix in prefixes
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix == prefixes
+
+    else:
+        prefixes = set(prefixes)
+
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix in prefixes
 
     return _func
 
@@ -134,10 +136,15 @@ def keep_object_prefixes(
 
 
 def _keep_object_prefixes_filter(prefixes: str | Iterable[str]) -> TriplePredicate[TripleType]:
-    prefixes = _cleanup_prefixes(prefixes)
+    if isinstance(prefixes, str):
 
-    def _func(triple: TripleType) -> bool:
-        return triple.object.prefix in prefixes
+        def _func(triple: TripleType) -> bool:
+            return triple.object.prefix == prefixes
+    else:
+        prefixes = set(prefixes)
+
+        def _func(triple: TripleType) -> bool:
+            return triple.object.prefix in prefixes
 
     return _func
 
@@ -168,10 +175,16 @@ def exclude_prefixes(
 
 
 def _exclude_prefixes_filter(prefixes: str | Iterable[str]) -> TriplePredicate[TripleType]:
-    prefixes = _cleanup_prefixes(prefixes)
+    if isinstance(prefixes, str):
 
-    def _func(triple: TripleType) -> bool:
-        return triple.subject.prefix not in prefixes and triple.object.prefix not in prefixes
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix != prefixes and triple.object.prefix != prefixes
+
+    else:
+        prefixes = set(prefixes)
+
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix not in prefixes and triple.object.prefix not in prefixes
 
     return _func
 
@@ -202,10 +215,16 @@ def exclude_subject_prefixes(
 
 
 def _exclude_subject_prefixes_filter(prefixes: str | Iterable[str]) -> TriplePredicate[TripleType]:
-    prefixes = _cleanup_prefixes(prefixes)
+    if isinstance(prefixes, str):
 
-    def _func(triple: TripleType) -> bool:
-        return triple.subject.prefix not in prefixes
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix != prefixes
+
+    else:
+        prefixes = set(prefixes)
+
+        def _func(triple: TripleType) -> bool:
+            return triple.subject.prefix not in prefixes
 
     return _func
 
@@ -236,10 +255,16 @@ def exclude_object_prefixes(
 
 
 def _exclude_object_prefixes_filter(prefixes: str | Iterable[str]) -> TriplePredicate[TripleType]:
-    prefixes = _cleanup_prefixes(prefixes)
+    if isinstance(prefixes, str):
 
-    def _func(triple: TripleType) -> bool:
-        return triple.object.prefix not in prefixes
+        def _func(triple: TripleType) -> bool:
+            return triple.object.prefix != prefixes
+
+    else:
+        prefixes = set(prefixes)
+
+        def _func(triple: TripleType) -> bool:
+            return triple.object.prefix not in prefixes
 
     return _func
 

tests/test_triples/test_filter.py

@@ -26,6 +26,7 @@
 m2 = Triple.from_curies(c2, exact_match.curie, c3)
 m3 = Triple.from_curies(c1, exact_match.curie, c3)
 m4 = Triple.from_curies(c1, subclass_of.curie, c4)
+M123 = [m1, m2, m3]
 converter = Converter.from_prefix_map(
     {
         "DOID": "http://purl.obolibrary.org/obo/DOID_",
@@ -46,84 +47,81 @@ class TestFilters(unittest.TestCase):
     def assert_triple_lists(self, expected: list[Triple], actual: Iterable[Triple]) -> None:
         """Test two triple lists are the same."""
         actual = list(actual)
-        self.assertEqual(
-            expected, list(actual), msg=f"\nExpected: {_x(expected)}\nActual: {_x(actual)}"
-        )
+        self.assertEqual(expected, actual, msg=f"\nExpected: {_x(expected)}\nActual: {_x(actual)}")
 
     def test_exclude_object_prefixes(self) -> None:
         """Test excluding object prefixes."""
-        self.assertEqual(
-            [m1],
-            list(exclude_object_prefixes([m1, m2, m3], {"umls"})),
-        )
-        self.assertEqual([m2, m3], list(exclude_object_prefixes([m1, m2, m3], {"mesh"})))
-        self.assertEqual([m1, m2, m3], list(exclude_object_prefixes([m1, m2, m3], {"DOID"})))
+        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.assertEqual([m1], list(exclude_prefixes([m1, m2, m3], {"umls"})))
-        self.assertEqual([m2], list(exclude_prefixes([m1, m2, m3], {"DOID"})))
-        self.assertEqual([m3], list(exclude_prefixes([m1, m2, m3], {"mesh"})))
+        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.assertEqual([m2], list(exclude_subject_prefixes([m1, m2, m3], {"DOID"})))
-        self.assertEqual([m1, m2, m3], list(exclude_subject_prefixes([m1, m2, m3], {"umls"})))
-        self.assertEqual([m1, m3], list(exclude_subject_prefixes([m1, m2, m3], {"mesh"})))
+        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"}))
 
     def test_exclude_same_prefixes(self) -> None:
         """Test excluding same prefixes."""
-        self.assertEqual([m1, m2, m3], list(exclude_same_prefixes([m1, m2, m3, m4])))
+        self.assert_triple_lists(M123, exclude_same_prefixes([m1, m2, m3, m4]))
 
     def test_exclude_triples(self) -> None:
         """Test excluding triples."""
-        self.assertEqual([m1, m2], list(exclude_triples([m1, m2, m3], m3)))
-        self.assertEqual([m1, m2], list(exclude_triples([m1, m2, m3], [m3])))
+        self.assert_triple_lists([m1, m2], exclude_triples(M123, m3))
+        self.assert_triple_lists([m1, m2], exclude_triples(M123, [m3]))
 
     def test_keep_object_prefixes(self) -> None:
         """Test keeping object prefixes."""
-        self.assertEqual([m2, m3], list(keep_object_prefixes([m1, m2, m3], {"umls"})))
+        self.assert_triple_lists([m2, m3], keep_object_prefixes(M123, {"umls"}))
 
     def test_keep_prefixes(self) -> None:
         """Test keeping prefixes."""
-        self.assertEqual([m1], list(keep_prefixes([m1, m2, m3], {"DOID", "mesh"})))
+        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"}))
 
     def test_keep_subject_prefixes(self) -> None:
         """Test keeping subject prefixes."""
-        self.assertEqual([m1, m3], list(keep_subject_prefixes([m1, m2, m3], {"DOID"})))
+        self.assert_triple_lists([m1, m3], keep_subject_prefixes(M123, {"DOID"}))
 
     def test_keep_triple_by_hash(self) -> None:
         """Test keeping triples by hash."""
-        self.assertEqual(
-            [m1], list(keep_triples_by_hash([m1, m2, m3], converter, converter.hash_triple(m1)))
+        self.assert_triple_lists(
+            [m1], keep_triples_by_hash(M123, converter, converter.hash_triple(m1))
         )
-        self.assertEqual(
+        self.assert_triple_lists(
             [m1, m2],
-            list(
-                keep_triples_by_hash(
-                    [m1, m2, m3], converter, [converter.hash_triple(m2), converter.hash_triple(m1)]
-                )
+            keep_triples_by_hash(
+                M123, converter, [converter.hash_triple(m2), converter.hash_triple(m1)]
             ),
         )
 
     def test_keep_references_either(self) -> None:
         """Test keeping references."""
-        self.assert_triple_lists([m1, m3], keep_references_either([m1, m2, m3], [r1]))
-        self.assert_triple_lists([m1, m2], keep_references_either([m1, m2, m3], [r2]))
-        self.assert_triple_lists([m2, m3], keep_references_either([m1, m2, m3], [r3]))
-        self.assert_triple_lists([m1, m2, m3], keep_references_either([m1, m2, m3], [r1, r2]))
-        self.assert_triple_lists([m1, m2, m3], keep_references_either([m1, m2, m3], [r2, r3]))
-        self.assert_triple_lists([m1, m2, m3], keep_references_either([m1, m2, m3], [r1, r2, r3]))
+        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]))
+        self.assert_triple_lists(M123, keep_references_either(M123, [r1, r2]))
+        self.assert_triple_lists(M123, keep_references_either(M123, [r2, r3]))
+        self.assert_triple_lists(M123, keep_references_either(M123, [r1, r2, r3]))
 
     def test_keep_references_both(self) -> None:
         """Test keeping references."""
-        self.assert_triple_lists([m1], keep_references_both([m1, m2, m3], [r1, r2]))
-        self.assert_triple_lists([m2], keep_references_both([m1, m2, m3], [r2, r3]))
-        self.assert_triple_lists([m3], keep_references_both([m1, m2, m3], [r1, r3]))
-        self.assert_triple_lists([m1, m2, m3], keep_references_both([m1, m2, m3], [r1, r2, r3]))
+        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.assertEqual([m2], list(exclude_references([m1, m2, m3], [r1])))
-        self.assertEqual([m3], list(exclude_references([m1, m2, m3], [r2])))
-        self.assertEqual([m1], list(exclude_references([m1, m2, m3], [r3])))
+        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]))