Represent triples with a Pydantic model (#172)

Author: cthoyt

src/curies/triples.py

@@ -4,14 +4,15 @@
 
 import csv
 import gzip
-from collections.abc import Generator, Iterable
+from collections.abc import Generator, Iterable, Sequence
 from contextlib import contextmanager
 from pathlib import Path
-from typing import NamedTuple, TextIO
+from typing import TextIO
 
+from pydantic import BaseModel
 from typing_extensions import Self
 
-from curies import Reference
+from .api import Reference
 
 __all__ = [
     "Triple",
@@ -20,24 +21,32 @@
 ]
 
 
-class Triple(NamedTuple):
-    """A three-tuple of reference, useful for semantic web applications."""
+class Triple(BaseModel):
+    """A model for a triple of subject-predicate-object triple."""
 
     subject: Reference
     predicate: Reference
     object: Reference
 
     @classmethod
-    def from_curies(cls, subject_curie: str, predicate_curie: str, object_curie: str) -> Self:
+    def from_curies(
+        cls,
+        subject_curie: str,
+        predicate_curie: str,
+        object_curie: str,
+        *,
+        reference_cls: type[Reference] = Reference,
+    ) -> Self:
         """Construct a triple from three CURIE strings."""
         return cls(
-            Reference.from_curie(subject_curie),
-            Reference.from_curie(predicate_curie),
-            Reference.from_curie(object_curie),
+            subject=reference_cls.from_curie(subject_curie),
+            predicate=reference_cls.from_curie(predicate_curie),
+            object=reference_cls.from_curie(object_curie),
         )
 
 
-HEADER = Triple._fields
+#: the default header for a three-column file representing triples
+HEADER = list(Triple.model_fields)
 
 
 @contextmanager
@@ -49,11 +58,15 @@ def _get_file(path: str | Path, read: bool) -> Generator[TextIO, None, None]:
         yield open(path, mode="r" if read else "w")
 
 
-def write_triples(triples: Iterable[Triple], path: str | Path) -> None:
+def write_triples(
+    triples: Iterable[Triple], path: str | Path, *, header: Sequence[str] | None = None
+) -> None:
     """Write triples to a file."""
+    if header is None:
+        header = HEADER
     with _get_file(path, read=False) as file:
         writer = csv.writer(file, delimiter="\t")
-        writer.writerow(HEADER)
+        writer.writerow(header)
         writer.writerows(
             (triple.subject.curie, triple.predicate.curie, triple.object.curie)
             for triple in triples
@@ -69,9 +82,9 @@ def read_triples(path: str | Path, *, reference_cls: type[Reference] | None = No
         _header = next(reader)
         return [
             Triple(
-                reference_cls.from_curie(subject_curie),
-                reference_cls.from_curie(predicate_curie),
-                reference_cls.from_curie(object_curie),
+                subject=reference_cls.from_curie(subject_curie),
+                predicate=reference_cls.from_curie(predicate_curie),
+                object=reference_cls.from_curie(object_curie),
             )
             for subject_curie, predicate_curie, object_curie in reader
         ]

tests/test_triples.py

@@ -1,5 +1,6 @@
 """Tests for triples."""
 
+import itertools as itt
 import tempfile
 import unittest
 from pathlib import Path
@@ -18,11 +19,13 @@ def test_roundtrip(self) -> None:
             Triple.from_curies("a:1", "a:2", "a:4"),
         ]
         with tempfile.TemporaryDirectory() as directory:
-            for path in [
+            paths = [
                 Path(directory).joinpath("test.tsv.gz"),
                 Path(directory).joinpath("test.tsv"),
-            ]:
+            ]
+            headers = [None, ("a", "b", "c")]
+            for path, header in itt.product(paths, headers):
                 with self.subTest(path=path):
-                    write_triples(triples, path)
+                    write_triples(triples, path, header=header)
                     reconstituted = read_triples(path)
                     self.assertEqual(triples, reconstituted)