Repository: cogent3/cogent3
Pull request #1010 (2022-02-03 04:53:22)

ENH: public decorator supporting deserialising custom objects
[NEW] the register_deserialiser class takes a series of strings
    that serve to uniquely identify the "type" value in a dict
    to be reconstituted using the decorated function.

    This enables support for user defined custom json storage.


Patch:
diff --git a/src/cogent3/util/deserialise.py b/src/cogent3/util/deserialise.py
index 95acbedf7..93f8490ee 100644
--- a/src/cogent3/util/deserialise.py
+++ b/src/cogent3/util/deserialise.py
@@ -20,6 +20,38 @@ __maintainer__ = "Gavin Huttley"
 __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Production"
 
+_deserialise_func_map = {}
+
+
+class register_deserialiser:
+    """
+    registration decorator for functions to inflate objects that were
+    serialised using json.
+
+    Functions are added to a dict which is used by the deserialise_object()
+    function. The type string(s) must uniquely identify the appropriate
+    value for the dict 'type' entry, e.g. 'cogent3.util.table.Table'.
+
+    Parameters
+    ----------
+    args: str or sequence of str
+        must be unique
+    """
+
+    def __init__(self, *args):
+        for type_str in args:
+            if not isinstance(type_str, str):
+                raise TypeError(f"{type_str!r} is not a string")
+            assert (
+                type_str not in _deserialise_func_map
+            ), f"{type_str!r} already in {list(_deserialise_func_map)}"
+        self._type_str = args
+
+    def __call__(self, func):
+        for type_str in self._type_str:
+            _deserialise_func_map[type_str] = func
+        return func
+
 
 def _get_class(provenance):
     index = provenance.rfind(".")
@@ -32,6 +64,11 @@ def _get_class(provenance):
     return klass
 
 
+@register_deserialiser(
+    "cogent3.util.table.Table",
+    "cogent3.util.dict_array.DictArray",
+    "cogent3.evolve.fast_distance.DistanceMatrix",
+)
 def deserialise_tabular(data):
     """deserialising DictArray, Table instances"""
     data.pop("version", None)
@@ -61,6 +98,7 @@ def deserialise_tabular(data):
     return result
 
 
+@register_deserialiser("cogent3.app.composable.NotCompleted")
 def deserialise_not_completed(data):
     """deserialising NotCompletedResult"""
     data.pop("version", None)
@@ -100,17 +138,14 @@ def deserialise_annotation(data, parent):
     parent.annotations += tuple(annots)
 
 
+@register_deserialiser("cogent3.app.result")
 def deserialise_result(data):
     """returns a result object"""
     data.pop("version", None)
     klass = _get_class(data.pop("type"))
     kwargs = data.pop("result_construction")
     result = klass(**kwargs)
-    if "items" in data:
-        items = data.pop("items")
-    else:
-        # retain support for the old style result serialisation
-        items = data.items()
+    items = data.pop("items") if "items" in data else data.items()
     for key, value in items:
         # only deserialise the result object, other attributes loaded as
         # required
@@ -123,6 +158,7 @@ def deserialise_result(data):
     return result
 
 
+@register_deserialiser("cogent3.core.moltype")
 def deserialise_moltype(data):
     """returns a cogent3 MolType instance, or a CodonAlphabet"""
     data.pop("version", None)
@@ -139,6 +175,7 @@ def deserialise_moltype(data):
     return result
 
 
+@register_deserialiser("cogent3.core.alphabet")
 def deserialise_alphabet(data):
     """returns a cogent3 Alphabet instance"""
     data.pop("version", None)
@@ -155,6 +192,7 @@ def deserialise_alphabet(data):
     return result
 
 
+@register_deserialiser("cogent3.core.sequence")
 def deserialise_seq(data, aligned=False):
     """deserialises sequence and any annotations
 
@@ -193,6 +231,7 @@ def deserialise_seq(data, aligned=False):
     return result
 
 
+@register_deserialiser("cogent3.core.alignment")
 def deserialise_seq_collections(data):
     """returns a cogent3 sequence/collection/alignment instance"""
     # We first try to load moltype/alphabet using get_moltype
@@ -219,6 +258,7 @@ def deserialise_seq_collections(data):
     return result
 
 
+@register_deserialiser("cogent3.core.tree")
 def deserialise_tree(data):
     """returns a cogent3 PhyloNode instance"""
     data.pop("version", None)
@@ -232,6 +272,9 @@ def deserialise_tree(data):
     return tree
 
 
+@register_deserialiser(
+    "cogent3.evolve.substitution_model", "cogent3.evolve.ns_substitution_model"
+)
 def deserialise_substitution_model(data):
     """returns a cogent3 substitution model instance"""
     from cogent3.evolve.models import get_model
@@ -254,6 +297,7 @@ def deserialise_substitution_model(data):
     return sm
 
 
+@register_deserialiser("cogent3.evolve.parameter_controller")
 def deserialise_likelihood_function(data):
     """returns a cogent3 likelihood function instance"""
     data.pop("version", None)
@@ -294,6 +338,11 @@ def deserialise_object(data):
     -------
     If the dict from json.loads does not contain a "type" key, the object will
     be returned as is. Otherwise, it will be deserialised to a cogent3 object.
+
+    Notes
+    -----
+    The value of the "type" key is used to identify the specific function for recreating
+    the original instance.
     """
     if path_exists(data):
         with open_(data) as infile:
@@ -306,33 +355,11 @@ def deserialise_object(data):
     if type_ is None:
         return data
 
-    if "core.sequence" in type_:
-        func = deserialise_seq
-    elif "core.alignment" in type_:
-        func = deserialise_seq_collections
-    elif "core.tree" in type_:
-        func = deserialise_tree
-    elif (
-        "evolve.substitution_model" in type_ or "evolve.ns_substitution_model" in type_
-    ):
-        func = deserialise_substitution_model
-    elif "evolve.parameter_controller" in type_:
-        func = deserialise_likelihood_function
-    elif "core.moltype" in type_:
-        func = deserialise_moltype
-    elif "core.alphabet" in type_:
-        func = deserialise_alphabet
-    elif "app.result" in type_:
-        func = deserialise_result
-    elif "notcompleted" in type_.lower():
-        func = deserialise_not_completed
-    elif type_.lower().endswith("table"):
-        func = deserialise_tabular
-    elif "dictarray" in type_.lower():
-        func = deserialise_tabular
-    elif "distancematrix" in type_.lower():
-        func = deserialise_tabular
+    for type_str, func in _deserialise_func_map.items():
+        if type_str in type_:
+            break
     else:
         msg = f"deserialising '{type_}' from json"
         raise NotImplementedError(msg)
+
     return func(data)


Test patch:
diff --git a/tests/test_util/test_deserialise.py b/tests/test_util/test_deserialise.py
index 3d6f222d4..299d2940b 100644
--- a/tests/test_util/test_deserialise.py
+++ b/tests/test_util/test_deserialise.py
@@ -483,6 +483,28 @@ class TestDeserialising(TestCase):
             str(got.defn_for["alignment"].assignments),
         )
 
+    def test_custom_deserialiser(self):
+        """correctly registers a function to inflate a custom object"""
+        from cogent3.util.deserialise import register_deserialiser
+
+        @register_deserialiser("myfunkydata")
+        def astuple(data):
+            data.pop("type")
+            return tuple(data["data"])
+
+        orig = {"type": "myfunkydata", "data": (1, 2, 3)}
+        txt = json.dumps(orig)
+        got = deserialise_object(txt)
+        self.assertEqual(got, (1, 2, 3))
+        self.assertIsInstance(got, tuple)
+
+        with self.assertRaises(TypeError):
+
+            @register_deserialiser
+            def astupled(data):
+                data.pop("type")
+                return tuple(data["data"])
+
 
 if __name__ == "__main__":
     main()