Repository: cogent3/cogent3
Pull request #1227 (2022-12-13 07:13:21)

ENH: support for reading and writing from a db backend


Patch:
diff --git a/src/cogent3/app/io_new.py b/src/cogent3/app/io_new.py
index 53ebe4bd5..d4c189fd9 100644
--- a/src/cogent3/app/io_new.py
+++ b/src/cogent3/app/io_new.py
@@ -1,6 +1,10 @@
+import contextlib
 import json
+import pickle
 
 from enum import Enum
+from gzip import compress as gzip_compress
+from gzip import decompress as gzip_decompress
 from typing import Optional, Union
 
 import numpy
@@ -45,6 +49,77 @@ __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Alpha"
 
 
+@define_app
+def pickle_it(data: SerialisableType) -> bytes:
+    return pickle.dumps(data)
+
+
+@define_app
+def unpickle_it(data: bytes) -> SerialisableType:
+    return pickle.loads(data)
+
+
+@define_app
+class compress:
+    def __init__(self, compressor: callable = gzip_compress):
+        """
+        Parameters
+        ----------
+        compressor
+            function for compressing bytes data, defaults to gzip
+        """
+        self.compressor = compressor
+
+    def main(self, data: bytes) -> bytes:
+        return self.compressor(data)
+
+
+@define_app
+class decompress:
+    def __init__(self, decompressor: callable = gzip_decompress):
+        """
+        Parameters
+        ----------
+        decompressor
+            a function for decompression, defaults to the gzip decompress
+            function
+        """
+        self.decompressor = decompressor
+
+    def main(self, data: bytes) -> bytes:
+        return self.decompressor(data)
+
+
+def _as_dict(obj) -> dict:
+    with contextlib.suppress(AttributeError):
+        obj = obj.to_rich_dict()
+    return obj
+
+
+@define_app
+class to_primitive:
+    """convert an object to primitive python types suitable for serialisation"""
+
+    def __init__(self, convertor: callable = _as_dict):
+        self.convertor = convertor
+
+    def main(self, data: SerialisableType) -> SerialisableType:
+        """either json convertor a dict from a cogent3 object"""
+        return self.convertor(data)
+
+
+@define_app
+class from_primitive:
+    """deserialises from primitive python types"""
+
+    def __init__(self, deserialiser: callable = deserialise_object):
+        self.deserialiser = deserialiser
+
+    def main(self, data: SerialisableType) -> SerialisableType:
+        """either json or a dict from a cogent3 object"""
+        return self.deserialiser(data)
+
+
 class tabular(Enum):
     table = "table"
     distances = "distances"
@@ -242,6 +317,28 @@ class load_json:
         return result
 
 
+@define_app(app_type=LOADER)
+class load_db:
+    """Loads serialised cogent3 objects from a db.
+    Returns whatever object type was stored."""
+
+    def __init__(self, deserialiser: callable = unpickle_it() + from_primitive()):
+        self.deserialiser = deserialiser
+
+    def main(self, identifier: IdentifierType) -> SerialisableType:
+        """returns deserialised object"""
+        data = identifier.read()
+        # do we need to inject identifier attribute?
+        result = self.deserialiser(data)
+        if hasattr(result, "info"):
+            result.info["source"] = result.info.get("source", identifier)
+        else:
+            with contextlib.suppress(AttributeError):
+                identifier = getattr(result, "source", identifier)
+                setattr(result, "source", identifier)
+        return result
+
+
 @define_app(app_type=WRITER)
 class write_json:
     def __init__(
@@ -305,3 +402,39 @@ class write_tabular:  # todo doctsring
 
         output = data.to_string(format=self._format)
         return self.data_store.write(unique_id=identifier, data=output)
+
+
+@define_app(app_type=WRITER)
+class write_db:
+    """Write serialised objects to a database instance."""
+
+    def __init__(
+        self,
+        data_store: DataStoreABC,
+        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
+        """
+        identifier = identifier or get_data_source(data)
+        data = self.serialiser(data)
+
+        if isinstance(data, NotCompleted):
+            return self.data_store.write_incomplete(identifier, data)
+
+        return self.data_store.write(unique_id=identifier, data=data)
diff --git a/src/cogent3/util/deserialise.py b/src/cogent3/util/deserialise.py
index d0676d129..f894dc0cc 100644
--- a/src/cogent3/util/deserialise.py
+++ b/src/cogent3/util/deserialise.py
@@ -358,7 +358,7 @@ def deserialise_object(data):
     if type(data) is str:
         data = json.loads(data)
 
-    type_ = data.get("type", None)
+    type_ = data.get("type", None) if hasattr(data, "get") else None
     if type_ is None:
         return data
 


Test patch:
diff --git a/tests/test_app/test_io_new.py b/tests/test_app/test_io_new.py
index a599cc78e..f6df14f87 100644
--- a/tests/test_app/test_io_new.py
+++ b/tests/test_app/test_io_new.py
@@ -1,8 +1,15 @@
+import bz2
+import gzip
+import json
+import pickle
+
 from pathlib import Path
 
 import numpy
 import pytest
 
+from numpy.testing import assert_allclose
+
 from cogent3 import DNA
 from cogent3.app import io_new as io_app
 from cogent3.app.composable import NotCompleted, source_proxy
@@ -12,6 +19,8 @@ from cogent3.core.profile import PSSM, MotifCountsArray, MotifFreqsArray
 from cogent3.evolve.fast_distance import DistanceMatrix
 from cogent3.maths.util import safe_log
 from cogent3.parse.sequence import PARSERS
+from cogent3.util.deserialise import deserialise_object
+from cogent3.util.table import Table
 
 
 __author__ = "Gavin Huttley"
@@ -23,10 +32,6 @@ __maintainer__ = "Gavin Huttley"
 __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Alpha"
 
-from numpy.testing import assert_allclose
-
-from cogent3.util.table import Table
-
 
 DATA_DIR = Path(__file__).parent.parent / "data"
 
@@ -367,5 +372,76 @@ def test_write_json_with_info(w_dir_dstore):
     assert got["dna"] == DNA
 
 
+@pytest.mark.parametrize(
+    "serialiser,deserialiser",
+    (
+        (json.dumps, json.loads),
+        (pickle.dumps, pickle.loads),
+        (lambda x: x, deserialise_object),
+    ),
+)
+def test_deserialiser(serialiser, deserialiser):
+    data = {"1": 1, "abc": [1, 2]}
+    deserialised = io_app.from_primitive(deserialiser=deserialiser)
+    assert deserialised(serialiser(data)) == data
+
+
+@pytest.mark.parametrize(
+    "data,dser", (([1, 2, 3], None), (DNA, io_app.from_primitive()))
+)
+def test_pickle_unpickle_apps(data, dser):
+    pkld = io_app.to_primitive() + io_app.pickle_it()
+    upkld = io_app.unpickle_it() + io_app.from_primitive()
+    # need to add custom deserialiser for cogent3 objects
+    # upkld = upkld if dser is None else upkld + dser
+    assert upkld(pkld(data)) == data
+
+
+def test_pickle_it_unpickleable():
+    def foo():  # can't pickle a local function
+        ...
+
+    app = io_app.pickle_it()
+    got = app(foo)
+    assert isinstance(got, NotCompleted)
+
+
+@pytest.mark.parametrize(
+    "compress,decompress",
+    ((bz2.compress, bz2.decompress), (gzip.compress, gzip.decompress)),
+)
+def test_compress_decompress(compress, decompress):
+    data = pickle.dumps({"1": 1, "abc": [1, 2]})
+
+    decompressor = io_app.decompress(decompressor=decompress)
+    compressor = io_app.compress(compressor=compress)
+
+    assert decompressor(compressor(data)) == data
+
+
+@pytest.mark.parametrize("data", ([1, 2, 3], DNA)[1:])
+def test_pickled_compress_roundtrip(data):
+    serialised = io_app.to_primitive() + io_app.pickle_it() + io_app.compress()
+    deserialised = io_app.decompress() + io_app.unpickle_it() + io_app.from_primitive()
+    s = serialised(data)
+    d = deserialised(s)
+    assert d == data
+
+
 # todo test objects where there is no unique_id provided, or inferrable,
 # should return NotCompletedError
+
+
+def test_write_db_load_db(fasta_dir, tmp_dir):
+    from cogent3.app.sqlite_data_store import DataStoreSqlite
+
+    orig_dstore = DataStoreDirectory(fasta_dir, suffix="fasta")
+    data_store = DataStoreSqlite(tmp_dir / "test.sqlitedb", mode="w")
+    load_seqs = io_app.load_unaligned()
+    writer = io_app.write_db(data_store=data_store)
+    reader = io_app.load_db()
+    for m in orig_dstore:
+        orig = load_seqs(m)
+        m = writer(orig, identifier=m.unique_id)
+        read = reader(m)
+        assert orig == read
diff --git a/tests/test_util/test_deserialise.py b/tests/test_util/test_deserialise.py
index c8989959e..f4ffb2bda 100644
--- a/tests/test_util/test_deserialise.py
+++ b/tests/test_util/test_deserialise.py
@@ -465,6 +465,9 @@ class TestDeserialising(TestCase):
         jdata = json.dumps(data)
         got = deserialise_object(jdata)
         self.assertEqual(got, data)
+        data = range(4)
+        got = deserialise_object(data)
+        assert got is data
 
     def test_deserialise_likelihood_function1(self):
         """correctly deserialise data into likelihood function"""