Add testing for graph I/O

src/semra/io/__init__.py

@@ -1,6 +1,6 @@
 """I/O functions for SeMRA."""
 
-from .graph import from_digraph, to_digraph, to_multidigraph
+from .graph import from_digraph, from_multidigraph, to_digraph, to_multidigraph
 from .io import (
     from_bioontologies,
     from_cache_df,
@@ -21,6 +21,7 @@
     "from_cache_df",
     "from_digraph",
     "from_jsonl",
+    "from_multidigraph",
     "from_pickle",
     "from_pyobo",
     "from_sssom",

src/semra/io/graph.py

@@ -14,6 +14,7 @@
     "DIGRAPH_DATA_KEY",
     "MULTIDIGRAPH_DATA_KEY",
     "from_digraph",
+    "from_multidigraph",
     "to_digraph",
     "to_multidigraph",
 ]
@@ -94,3 +95,11 @@ def to_multidigraph(mappings: t.Iterable[Mapping], *, progress: bool = False) ->
             **{MULTIDIGRAPH_DATA_KEY: mapping.evidence},
         )
     return graph
+
+
+def from_multidigraph(graph: nx.MultiDiGraph) -> list[Mapping]:
+    """Extract mappings from a multi-directed graph data model."""
+    return [
+        Mapping(s=s, p=p, o=o, evidence=data[MULTIDIGRAPH_DATA_KEY])
+        for s, o, p, data in graph.edges(keys=True, data=True)
+    ]

src/semra/struct.py

@@ -359,6 +359,9 @@ def key(self) -> StrTriple:
         """Get a hashable key for the mapping, based on the subject, predicate, and object."""
         return triple_key(self.triple)
 
+    def __lt__(self, other: Mapping) -> bool:
+        return self.triple < other.triple
+
     @classmethod
     def from_triple(cls, triple: Triple, evidence: list[Evidence] | None = None) -> Mapping:
         """Instantiate a mapping from a triple."""

tests/test_io.py

@@ -9,7 +9,16 @@
 import pandas as pd
 
 from semra import Mapping, MappingSet, ReasonedEvidence, Reference, SimpleEvidence
-from semra.io import from_jsonl, from_pyobo, from_sssom_df, write_jsonl
+from semra.io import (
+    from_digraph,
+    from_jsonl,
+    from_multidigraph,
+    from_pyobo,
+    from_sssom_df,
+    to_digraph,
+    to_multidigraph,
+    write_jsonl,
+)
 from semra.rules import (
     BEN_ORCID,
     CHAIN_MAPPING,
@@ -44,7 +53,7 @@ def test_from_pyobo(self) -> None:
             self.assertEqual("mesh", mapping.o.prefix)
 
 
-class TestIO(unittest.TestCase):
+class TestSSSOM(unittest.TestCase):
     """Tests for I/O functions."""
 
     def test_from_sssom_df(self) -> None:
@@ -173,8 +182,12 @@ def test_from_sssom_df_with_license(self) -> None:
         )
         self.assertEqual(expected_mappings, actual_mappings)
 
-    def test_jsonl(self) -> None:
-        """Test JSONL I/O."""
+
+class TestIO(unittest.TestCase):
+    """Test I/O funcitons."""
+
+    def setUp(self) -> None:
+        """Set up the test case."""
         r1 = Reference.from_curie("mesh:C406527", name="R 115866")
         r2 = Reference.from_curie("chebi:101854", name="talarozole")
         r3 = Reference.from_curie("chembl.compound:CHEMBL459505", name="TALAROZOLE")
@@ -214,13 +227,25 @@ def test_jsonl(self) -> None:
         m3_e1 = ReasonedEvidence(justification=CHAIN_MAPPING, mappings=[m1, m2])
         m3 = Mapping.from_triple(t3, evidence=[m3_e1])
 
-        mappings = [m1, m2, m3]
+        self.mappings = [m1, m2, m3]
 
+    def test_jsonl(self) -> None:
+        """Test JSONL I/O."""
         with tempfile.TemporaryDirectory() as directory_:
             for path in [
                 Path(directory_).joinpath("test.jsonl"),
                 Path(directory_).joinpath("test.jsonl.gz"),
             ]:
-                write_jsonl(mappings, path)
+                write_jsonl(self.mappings, path)
                 new_mappings = from_jsonl(path, show_progress=False)
-                self.assertEqual(mappings, new_mappings)
+                self.assertEqual(self.mappings, new_mappings)
+
+    def test_digraph(self) -> None:
+        """Test I/O to a directed graph."""
+        self.assertEqual(sorted(self.mappings), sorted(from_digraph(to_digraph(self.mappings))))
+
+    def test_multidigraph(self) -> None:
+        """Test I/O with multi-directed graph."""
+        self.assertEqual(
+            sorted(self.mappings), sorted(from_multidigraph(to_multidigraph(self.mappings)))
+        )