Repository: cogent3/cogent3
Pull request #989 (2021-12-01 02:25:31)

MAINT: refactor code for inferring data source
[CHANGED] made public get_data_source() function under cogent3.app.data_store,
    which handles a non-exhaustive series of cases and is relied on by the
    data store classes for extracting the provenance of data objects
[CHANGED] the make_(un)aligned_seqs functions now inspect both info and source
    arguments for a source value. The source argument takes precedence.


Patch:
diff --git a/src/cogent3/__init__.py b/src/cogent3/__init__.py
index 4741b7925..1389924d3 100644
--- a/src/cogent3/__init__.py
+++ b/src/cogent3/__init__.py
@@ -150,7 +150,7 @@ def make_unaligned_seqs(
         other_kw = kw.pop(other_kw, None) or {}
         kw.update(other_kw)
     assert isinstance(info, dict), "info must be a dict"
-    info["source"] = source or "unknown"
+    info["source"] = source or info.get("source", "unknown")
 
     return SequenceCollection(
         data=data, moltype=moltype, label_to_name=label_to_name, info=info, **kw
@@ -193,7 +193,7 @@ def make_aligned_seqs(
         other_kw = kw.pop(other_kw, None) or {}
         kw.update(other_kw)
     assert isinstance(info, dict), "info must be a dict"
-    info["source"] = source or "unknown"
+    info["source"] = source or info.get("source", "unknown")
     klass = ArrayAlignment if array_align else Alignment
     return klass(
         data=data, moltype=moltype, label_to_name=label_to_name, info=info, **kw
@@ -467,9 +467,9 @@ def load_table(
     if file_format == "json":
         return load_from_json(filename, (_Table,))
     elif file_format in ("pickle", "pkl"):
-        f = open_(filename, mode="rb")
-        loaded_table = pickle.load(f)
-        f.close()
+        with open_(filename, mode="rb") as f:
+            loaded_table = pickle.load(f)
+
         r = _Table()
         r.__setstate__(loaded_table)
         return r
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index bb782dd6d..5fa174144 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -26,6 +26,7 @@ from .data_store import (
     DataStoreMember,
     SingleReadDataStore,
     WritableDirectoryDataStore,
+    get_data_source,
 )
 
 
@@ -49,21 +50,6 @@ def _make_logfile_name(process):
     return result
 
 
-def _get_source(source):
-    if isinstance(source, str):
-        return str(source)
-
-    # todo maybe a dict? see about getting keys
-    try:
-        result = source.source
-    except AttributeError:
-        try:
-            result = source.info.source
-        except AttributeError:
-            result = None
-    return result
-
-
 def _get_origin(origin):
     return origin if type(origin) == str else origin.__class__.__name__
 
@@ -87,7 +73,7 @@ class NotCompleted(int):
         # todo this approach to caching persistent arguments for reconstruction
         # is fragile. Need an inspect module based approach
         origin = _get_origin(origin)
-        source = _get_source(source)
+        source = get_data_source(source)
         d = locals()
         d = {k: v for k, v in d.items() if k != "cls"}
         result = int.__new__(cls, False)
diff --git a/src/cogent3/app/data_store.py b/src/cogent3/app/data_store.py
index 384b93d91..36fd19105 100644
--- a/src/cogent3/app/data_store.py
+++ b/src/cogent3/app/data_store.py
@@ -45,6 +45,29 @@ RAISE = "raise"
 IGNORE = "ignore"
 
 
+def get_data_source(data) -> str:
+    """identifies attribute of data named 'source'
+
+    Notes
+    -----
+    Alignment objects have a source element in their info dict
+    """
+    if isinstance(data, (str, pathlib.Path)):
+        return str(data)
+
+    if hasattr(data, "source"):
+        return str(data.source)
+
+    if hasattr(data, "info"):
+        return get_data_source(data.info)
+
+    if isinstance(data, dict):
+        value = data.get("source")
+        return str(value) if value else None
+
+    return None
+
+
 def make_record_for_json(identifier, data, completed):
     """returns a dict for storage as json"""
     try:
@@ -407,12 +430,11 @@ class WritableDataStoreBase:
 
     def make_relative_identifier(self, data):
         """returns identifier for a new member relative to source"""
-        from cogent3.app.composable import _get_source
 
         if isinstance(data, DataStoreMember):
             data = data.name
         elif type(data) != str:
-            data = _get_source(data)
+            data = get_data_source(data)
             if data is None:
                 raise ValueError(
                     "objects for storage require either a "
diff --git a/src/cogent3/app/io.py b/src/cogent3/app/io.py
index 7de29a187..73e0e5fe0 100644
--- a/src/cogent3/app/io.py
+++ b/src/cogent3/app/io.py
@@ -32,7 +32,6 @@ from .composable import (
     ComposableTabular,
     NotCompleted,
     _checkpointable,
-    _get_source,
 )
 from .data_store import (
     IGNORE,
@@ -44,6 +43,7 @@ from .data_store import (
     ReadOnlyZippedDataStore,
     SingleReadDataStore,
     WritableTinyDbDataStore,
+    get_data_source,
     load_record_from_json,
     make_record_for_json,
 )
@@ -594,7 +594,7 @@ class write_db(_checkpointable):
         -------
         identifier
         """
-        data_source = _get_source(data)  # todo -- this is being ignored!
+        data_source = get_data_source(data)  # todo -- this is being ignored!
         if (data_source and identifier is not None) and str(data_source) != str(
             identifier
         ):


Test patch:
diff --git a/tests/test_app/test_data_store.py b/tests/test_app/test_data_store.py
index 30e722629..a91887291 100644
--- a/tests/test_app/test_data_store.py
+++ b/tests/test_app/test_data_store.py
@@ -20,6 +20,7 @@ from cogent3.app.data_store import (
     SingleReadDataStore,
     WritableDirectoryDataStore,
     WritableTinyDbDataStore,
+    get_data_source,
     load_record_from_json,
 )
 from cogent3.parse.fasta import MinimalFastaParser
@@ -854,6 +855,53 @@ class TestFunctions(TestCase):
             self.assertEqual(data_, expected)
             self.assertEqual(compl, True)
 
+    def test_get_data_source_str_pathlib(self):
+        """handles case where input is string object or pathlib object"""
+        for val_klass in (str, pathlib.Path):
+            value = val_klass("some/path.txt")
+            got = get_data_source(value)
+            self.assertEqual(got, str(value))
+
+    def test_get_data_source_seqcoll(self):
+        """handles case where input is sequence collection object"""
+        from cogent3 import make_unaligned_seqs
+
+        for val_klass in (str, pathlib.Path):
+            value = val_klass("some/path.txt")
+            obj = make_unaligned_seqs(data=dict(seq1="ACGG"), info=dict(source=value))
+            got = get_data_source(obj)
+            self.assertEqual(got, str(value))
+
+    def test_get_data_source_attr(self):
+        """handles case where input has source attribute string object or pathlib object"""
+
+        class dummy:
+            source = None
+
+        for val_klass in (str, pathlib.Path):
+            obj = dummy()
+            value = val_klass("some/path.txt")
+            obj.source = value
+            got = get_data_source(obj)
+            self.assertEqual(got, str(value))
+
+    def test_get_data_source_dict(self):
+        """handles case where input is dict (sub)class instance with top level source key"""
+        from cogent3.util.union_dict import UnionDict
+
+        for klass in (dict, UnionDict):
+            for val_klass in (str, pathlib.Path):
+                value = val_klass("some/path.txt")
+                data = klass(source=value)
+                got = get_data_source(data)
+                self.assertEqual(got, str(value))
+
+    def test_get_data_source_none(self):
+        """handles case where input does not have a source attribute or key"""
+        for data in (None, dict(), set(), dict(info=dict())):
+            got = get_data_source(data)
+            self.assertIsNone(got)
+
 
 if __name__ == "__main__":
     main()