Repository: cogent3/cogent3
Pull request #1269 (2023-01-16 05:14:47)

ENH: add argument to all writer apps for a naming callback
[NEW] id_from_source argument should be a callable that
    generates a unique ID from input data

[NEW] docstrings for the writers


Patch:
diff --git a/src/cogent3/app/io_new.py b/src/cogent3/app/io_new.py
index 581d5c7b2..20c9a053d 100644
--- a/src/cogent3/app/io_new.py
+++ b/src/cogent3/app/io_new.py
@@ -22,6 +22,7 @@ from cogent3.evolve.fast_distance import DistanceMatrix
 from cogent3.format.alignment import FORMATTERS
 from cogent3.parse.sequence import PARSERS
 from cogent3.util.deserialise import deserialise_object
+from cogent3.util.misc import extend_docstring_from
 from cogent3.util.table import Table
 
 from .composable import LOADER, WRITER, NotCompleted, define_app
@@ -466,14 +467,26 @@ class write_json:
     def __init__(
         self,
         data_store: DataStoreABC,
+        id_from_source: callable = get_unique_id,
     ):
+        """
+        Parameters
+        ----------
+        data_store
+            A writeable data store.
+        id_from_source
+            A function for creating a unique identifier based on the data
+            source. The default function removes path information and
+            filename + compression suffixes.
+        """
         self.data_store = data_store
         self._format = "json"
+        self._id_from_source = id_from_source
 
     def main(
-        self, data: SerialisableType, identifier: Optional[str] = None
+        self, /, data: SerialisableType, *, identifier: Optional[str] = None
     ) -> IdentifierType:
-        identifier = identifier or get_unique_id(data)
+        identifier = identifier or self._id_from_source(data)
         if isinstance(data, NotCompleted):
             return self.data_store.write_not_completed(
                 unique_id=f"{identifier}.json", data=data.to_json()
@@ -485,19 +498,26 @@ class write_json:
 
 
 @define_app(app_type=WRITER)
-class write_seqs:  # todo docstring
+class write_seqs:
+    @extend_docstring_from(write_json.__init__, pre=False)
     def __init__(
         self,
         data_store: DataStoreABC,
-        format="fasta",
+        id_from_source: callable = get_unique_id,
+        format: str = "fasta",
     ):
+        """
+        format
+            sequence format
+        """
         self.data_store = data_store
         self._formatter = FORMATTERS[format]
+        self._id_from_source = id_from_source
 
     def main(
-        self, data: SeqsCollectionType, identifier: Optional[str] = None
+        self, /, data: SeqsCollectionType, *, identifier: Optional[str] = None
     ) -> IdentifierType:
-        identifier = identifier or get_unique_id(data)
+        identifier = identifier or self._id_from_source(data)
         if isinstance(data, NotCompleted):
             return self.data_store.write_not_completed(
                 unique_id=f"{identifier}.json", data=data.to_json()
@@ -508,15 +528,26 @@ class write_seqs:  # todo docstring
 
 
 @define_app(app_type=WRITER)
-class write_tabular:  # todo doctsring
-    def __init__(self, data_store: DataStoreABC, format="tsv"):
+class write_tabular:
+    @extend_docstring_from(write_json.__init__, pre=False)
+    def __init__(
+        self,
+        data_store: DataStoreABC,
+        id_from_source: callable = get_unique_id,
+        format: str = "tsv",
+    ):
+        """
+        format
+            tabular format, e.g. 'csv' or 'tsv'
+        """
         self.data_store = data_store
+        self._id_from_source = id_from_source
         self._format = format
 
     def main(
-        self, data: TabularType, identifier: Optional[str] = None
+        self, /, data: TabularType, *, identifier: Optional[str] = None
     ) -> IdentifierType:
-        identifier = identifier or get_unique_id(data)
+        identifier = identifier or self._id_from_source(data)
         if isinstance(data, NotCompleted):
             return self.data_store.write_not_completed(
                 unique_id=f"{identifier}.json", data=data.to_json()
@@ -530,31 +561,27 @@ class write_tabular:  # todo doctsring
 class write_db:
     """Write serialised objects to a database instance."""
 
+    @extend_docstring_from(write_json.__init__, pre=False)
     def __init__(
         self,
         data_store: DataStoreABC,
+        id_from_source: callable = get_unique_id,
         serialiser: callable = to_primitive() + pickle_it(),
     ):
-        self.data_store = data_store
-        self.serialiser = serialiser
-
-    T = Union[SerialisableType, IdentifierType]
-
-    def main(self, /, data: SerialisableType, *, identifier=None) -> T:
         """
-        Parameters
-        ----------
-        data
-            object that has a `to_json()` method, or can be json serialised
-        identifier : str
-            if not provided, taken from data.source or data.info.source
-
-        Returns
-        -------
-        identifier
+        serialiser
+            A callable for serialising input data. By default, it converts
+            data into primitive python data types, pickling the result.
         """
-        identifier = identifier or get_unique_id(data)
-        blob = self.serialiser(data)
+        self.data_store = data_store
+        self._serialiser = serialiser
+        self._id_from_source = id_from_source
+
+    def main(
+        self, /, data: SerialisableType, *, identifier: Optional[str] = None
+    ) -> IdentifierType:
+        identifier = identifier or self._id_from_source(data)
+        blob = self._serialiser(data)
 
         if isinstance(data, NotCompleted):
             return self.data_store.write_not_completed(unique_id=identifier, data=blob)


Test patch:
diff --git a/tests/test_app/test_io_new.py b/tests/test_app/test_io_new.py
index ac1d7e746..78398aab8 100644
--- a/tests/test_app/test_io_new.py
+++ b/tests/test_app/test_io_new.py
@@ -2,6 +2,7 @@ import bz2
 import gzip
 import json
 import os
+import pathlib
 import pickle
 
 from pathlib import Path
@@ -87,7 +88,7 @@ def test_write_seqs(fasta_dir, tmp_dir):
         tmp_dir / "test_write_seqs", mode=Mode.w, suffix="fasta"
     )
     writer = io_app.write_seqs(out_data_store, format="fasta")
-    wrote = writer(seqs[0], datamember.unique_id)
+    wrote = writer(seqs[0], identifier=datamember.unique_id)
     assert isinstance(wrote, DataMember)
 
 
@@ -507,3 +508,54 @@ def test_define_data_store(fasta_dir):
 
     with pytest.raises(ValueError):
         _ = open_data_store(fasta_dir, 1)
+
+
+def seqs():
+    from cogent3 import make_unaligned_seqs
+
+    return make_unaligned_seqs(
+        data=dict(a="ACGG", b="GGC"), info=dict(source="dummy/blah.1.2.fa")
+    )
+
+
+def table():
+    from cogent3 import make_table
+
+    table = make_table(
+        data=dict(a=[0, 1, 2], b=[0, 1, 2]),
+    )
+    table.source = "dummy/blah.1.2.fa"
+    return table
+
+
+def dir_dstore(tmp_dir):
+    return open_data_store(tmp_dir / "sample", suffix="fa", mode="w")
+
+
+def db_dstore(tmp_dir):
+    return open_data_store(tmp_dir / "sample.sqlitedb", mode="w")
+
+
+@pytest.mark.parametrize(
+    "writer,data,dstore",
+    (
+        ("write_seqs", seqs(), dir_dstore),
+        ("write_db", seqs(), db_dstore),
+        ("write_json", seqs(), dir_dstore),
+        ("write_tabular", table(), dir_dstore),
+    ),
+)
+def test_writer_unique_id_arg(tmp_dir, writer, data, dstore):
+    def uniqid(source):
+        from cogent3.app.data_store_new import get_data_source
+
+        name = pathlib.Path(get_data_source(source)).name
+        return name.split(".")[0]
+
+    writer = get_app(writer, data_store=dstore(tmp_dir), id_from_source=uniqid)
+    m = writer(data)
+    # directory data stores have a suffix, so we create expected name using that
+    suffix = getattr(writer.data_store, "suffix", "")
+    suffix = f".{suffix}" if suffix else suffix
+    expect = f"blah{suffix}"
+    assert m.unique_id == expect