Repository: cogent3/cogent3
Pull request #1041 (2022-05-04 03:25:32)

MAINT: improve consistency of DataStore write methods


Patch:
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index 396e52b0b..279f35be9 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -474,8 +474,7 @@ class Composable(ComposableType):
                 # Note: we directly call .write() so NotCompleted's don't
                 # get blocked from being written by __call__()
                 outcome = self.write(data=result)
-                assert isinstance(outcome, DataStoreMember)
-                if result:
+                if result and isinstance(outcome, DataStoreMember):
                     input_id = outcome.name
                 else:
                     input_id = get_data_source(result)
@@ -678,7 +677,7 @@ class _checkpointable(Composable):
             exists = False
         return exists
 
-    def write(self, data):
+    def write(self, data) -> DataStoreMember:
         # over-ride in subclass
         raise NotImplementedError
 
diff --git a/src/cogent3/app/data_store.py b/src/cogent3/app/data_store.py
index 516b76f0f..fedb86174 100644
--- a/src/cogent3/app/data_store.py
+++ b/src/cogent3/app/data_store.py
@@ -503,7 +503,7 @@ class WritableDataStoreBase:
 
         return relativeid
 
-    def write_incomplete(self, identifier, not_completed):
+    def write_incomplete(self, identifier, not_completed) -> DataStoreMember:
         """
 
         Parameters
@@ -524,9 +524,9 @@ class WritableDataStoreBase:
 
         record = make_record_for_json(identifier, not_completed, False)
         record = json.dumps(record)
-        self.write(identifier, record)
+        return self.write(identifier, record)
 
-    def write(self, identifier, data, *args, **kwargs):
+    def write(self, identifier, data, *args, **kwargs) -> DataStoreMember:
         """
         Parameters
         ----------
diff --git a/src/cogent3/app/io.py b/src/cogent3/app/io.py
index 158d6c67e..3d17a57ce 100644
--- a/src/cogent3/app/io.py
+++ b/src/cogent3/app/io.py
@@ -428,6 +428,11 @@ class write_tabular(_checkpointable, ComposableTabular):
         self._format = format
 
     def write(self, data, identifier=None):
+        from cogent3.app.composable import NotCompleted
+
+        if isinstance(data, NotCompleted):
+            return self.data_store.write_incomplete(identifier, data)
+
         if not data:
             msg = f"{self.__class__.__name__!r} does not support writing {data!r}"
             raise NotImplementedError(msg)
@@ -436,8 +441,8 @@ class write_tabular(_checkpointable, ComposableTabular):
             identifier = self._make_output_identifier(data)
 
         output = data.to_string(format=self._format)
-        self.data_store.write(identifier, output)
-        return identifier
+        stored = self.data_store.write(identifier, output)
+        return stored
 
 
 class write_seqs(_checkpointable):
@@ -495,6 +500,11 @@ class write_seqs(_checkpointable):
         self._load_checkpoint = loader
 
     def write(self, data, identifier=None):
+        from cogent3.app.composable import NotCompleted
+
+        if isinstance(data, NotCompleted):
+            return self.data_store.write_incomplete(identifier, data)
+
         if not data:
             msg = f"{self.__class__.__name__!r} does not support writing {data!r}"
             raise NotImplementedError(msg)
@@ -502,8 +512,16 @@ class write_seqs(_checkpointable):
         if identifier is None:
             identifier = self._make_output_identifier(data)
 
-        data.info.stored = self.data_store.write(identifier, data.to_fasta())
-        return identifier
+        stored = self.data_store.write(identifier, self._formatter(data.to_dict()))
+        if hasattr(data, "info"):
+            data.info["stored"] = stored
+        else:
+            try:
+                data.stored = stored
+            except AttributeError:
+                pass
+
+        return stored
 
 
 class load_json(Composable):
@@ -582,7 +600,7 @@ class write_json(_checkpointable):
         out = make_record_for_json(os.path.basename(identifier), data, True)
         out = json.dumps(out)
         stored = self.data_store.write(identifier, out)
-        # todo is anything actually using this stored attriubte? if not, delete this
+        # todo is anything actually using this stored attribute? if not, delete this
         #  code and all other cases
         if hasattr(data, "info"):
             data.info["stored"] = stored
@@ -591,7 +609,7 @@ class write_json(_checkpointable):
                 data.stored = stored
             except AttributeError:
                 pass
-        return identifier
+        return stored
 
 
 class load_db(Composable):


Test patch:
diff --git a/tests/test_app/test_composable.py b/tests/test_app/test_composable.py
index 19365746b..10c35d0da 100644
--- a/tests/test_app/test_composable.py
+++ b/tests/test_app/test_composable.py
@@ -46,7 +46,7 @@ class TestCheckpoint(TestCase):
 
             read_write = reader + writer
             got = read_write(path)  # should skip reading and return path
-            self.assertEqual(got, outpath)
+            self.assertTrue(got.endswith(outpath))
             read_write.disconnect()  # allows us to reuse bits
             read_write_degen = reader + writer + omit_degens
             # should return an alignment instance
diff --git a/tests/test_app/test_data_store.py b/tests/test_app/test_data_store.py
index ec8f0ac73..58d6c346e 100644
--- a/tests/test_app/test_data_store.py
+++ b/tests/test_app/test_data_store.py
@@ -645,7 +645,8 @@ class TinyDBDataStoreTests(TestCase):
             path = os.path.join(dirname, self.basedir)
             dstore = self.WriteClass(path, if_exists="overwrite")
             id_ = dstore.make_relative_identifier(incomplete[0])
-            dstore.write(id_, incomplete[1])
+            got = dstore.write(id_, incomplete[1])
+            self.assertIsInstance(got, DataStoreMember)
             for k in keys:
                 id_ = dstore.make_relative_identifier(k)
                 dstore.write(id_, self.data[k])
@@ -807,7 +808,9 @@ class TinyDBDataStoreTests(TestCase):
             path = dirname / f"{self.basedir}.tinydb"
             dstore = self.WriteClass(path, create=True)
             with self.assertRaises(ValueError):
-                dstore.write("1", dict(a=24, b="some text"))
+                got = dstore.write("1", dict(a=24, b="some text"))
+                # validate return type
+                self.assertIsInstance(got, DataStoreMember)
 
             dstore.write("1.json", dict(a=24, b="some text"))
             dstore.close()
diff --git a/tests/test_app/test_io.py b/tests/test_app/test_io.py
index 214584944..70ecf4749 100644
--- a/tests/test_app/test_io.py
+++ b/tests/test_app/test_io.py
@@ -16,6 +16,7 @@ from cogent3 import DNA
 from cogent3.app import align as align_app
 from cogent3.app import io as io_app
 from cogent3.app.composable import NotCompleted
+from cogent3.app.data_store import DataStoreMember
 from cogent3.app.io import write_db
 from cogent3.app.result import generic_result
 from cogent3.core.alignment import ArrayAlignment, SequenceCollection
@@ -23,7 +24,6 @@ from cogent3.core.profile import PSSM, MotifCountsArray, MotifFreqsArray
 from cogent3.evolve.fast_distance import DistanceMatrix
 from cogent3.maths.util import safe_log
 from cogent3.util.table import Table
-from cogent3.util.union_dict import UnionDict
 
 
 __author__ = "Gavin Huttley"
@@ -174,6 +174,7 @@ class TestIo(TestCase):
         with TemporaryDirectory(dir=".") as dirname:
             writer = io_app.write_seqs(dirname, if_exists="ignore")
             wrote = list(map(writer, alns))
+            self.assertIsInstance(wrote[0], DataStoreMember)
             written = list(io_app.findall(dirname, suffix="fasta"))
             for i, wrote in enumerate(written):
                 self.assertEqual(alns[i].info.stored, join(dirname, wrote))
@@ -216,6 +217,7 @@ class TestIo(TestCase):
             writer = write_db(outpath, create=True, if_exists="ignore")
             gr = _get_generic_result(join("blah", "delme.json"))
             got = writer(gr)
+            self.assertIsInstance(got, DataStoreMember)
             writer.data_store.db.close()
             dstore = io_app.get_data_store(f"{outpath}.tinydb", suffix="json")
             reader = io_app.load_db()
@@ -289,7 +291,8 @@ class TestIo(TestCase):
         with TemporaryDirectory(dir=".") as dirname:
             writer = io_app.write_tabular(data_path=dirname, format="tsv")
             outpath = join(dirname, "delme.tsv")
-            writer.write(mca, identifier=outpath)
+            got = writer.write(mca, identifier=outpath)
+            self.assertIsInstance(got, DataStoreMember)
             new = loader(outpath)
             # when written to file in tabular form
             # the loaded table will have dim-1 dim-2 as column labels
@@ -468,7 +471,8 @@ class TestIo(TestCase):
             obj = generic_result(source=join("blah", "delme.json"))
             obj["dna"] = DNA
             writer = io_app.write_json(outdir, create=True)
-            _ = writer(obj)
+            got = writer(obj)
+            self.assertIsInstance(got, DataStoreMember)
             reader = io_app.load_json()
             got = reader(join(outdir, "delme.json"))
             got.deserialised_values()