Repository: cogent3/cogent3
Pull request #2163 (2024-12-11 06:29:55)

ENH: consistent API to ease migration to new_type
## Summary by Sourcery

Enhance serialization methods with a 'for_pickle' parameter and implement deep copy functionality. Introduce new class methods to support migration to new types in alignment modules. Add serialization methods to 'MolType' class and corresponding tests for new features.

New Features:
- Introduce 'from_map_and_seq' and 'from_map_and_aligned_data_view' class methods in 'new_alignment' and 'alignment' modules to facilitate migration to new types.
- Add 'to_rich_dict' and 'to_json' methods to 'MolType' class for serialization.

Enhancements:
- Add a 'for_pickle' parameter to 'to_rich_dict' methods to control inclusion of type and version information.
- Implement '__deepcopy__' method for deep copying objects using 'from_rich_dict'.

Tests:
- Add tests for 'from_map_and_seq' and 'from_map_and_aligned_data_view' methods in 'new_alignment'.
- Add parameterized test for JSON roundtrip of different moltypes to ensure serialization and deserialization consistency.


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index 6fe5d768b..02b4d644d 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -2520,6 +2520,22 @@ class Aligned:
         """returns True if sequence has any annotations"""
         return self.data.is_annotated()
 
+    @classmethod
+    def from_map_and_seq(cls, indel_map, seq):  # pragma: no cover
+        # method for forwards compatability
+        # no coverage due to class deprecation
+        return cls(indel_map, seq)
+
+    @classmethod
+    def from_map_and_aligned_data_view(
+        cls,
+        indel_map,
+        seq,
+    ):  # pragma: no cover
+        # method for forwards compatability
+        # no coverage due to class deprecation
+        return cls(indel_map, seq)
+
 
 class AlignmentI:
     """Alignment interface object. Contains methods shared by implementations.
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index f76781343..05335f792 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -2795,6 +2795,41 @@ class Aligned:
         self._name = name or data.seqid
         self._annotation_db = annotation_db
 
+    @classmethod
+    def from_map_and_seq(cls, indel_map: IndelMap, seq: new_sequence.Sequence):
+        moltype = seq.moltype
+        # refactor: design
+        # this is a temporary approach during migration to new_types
+        # to support the sequence alignment algorithms
+        # a better solution is to create a AlignedDataView instance the
+        # map and seq directly without requiring a parent AlignedSeqsData
+        asd = AlignedSeqsData.from_seqs_and_gaps(
+            seqs={seq.name: numpy.array(seq)},
+            gaps={seq.name: indel_map.array},
+            alphabet=moltype.most_degen_alphabet(),
+        )
+
+        return cls(asd.get_view(seq.name), moltype)
+
+    @classmethod
+    def from_map_and_aligned_data_view(
+        cls,
+        indel_map: IndelMap,
+        seq: AlignedDataViewABC,
+    ):
+        moltype = seq.alphabet.moltype
+        seqid = seq.seqid
+        seq = seq.array_value
+        # refactor: design
+        # see above comment in from_map_and_seq
+        asd = AlignedSeqsData.from_seqs_and_gaps(
+            seqs={seqid: seq},
+            gaps={seqid: indel_map.array},
+            alphabet=moltype.most_degen_alphabet(),
+        )
+
+        return cls(asd.get_view(seqid), moltype)
+
     def __len__(self) -> int:
         return len(self.map)
 
@@ -6445,6 +6480,10 @@ class Alignment(SequenceCollection):
         # TODO remove when new_type is the default
         return self
 
+    def is_ragged(self) -> bool:
+        """by definition False for an Alignment"""
+        return False
+
 
 @register_deserialiser(get_object_provenance(Alignment))
 def deserialise_alignment(data: dict[str, str | dict[str, str]]) -> Alignment:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index 5a542671a..f59198dee 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
 
 import numba
 import numpy
+import typing_extensions
 
 from cogent3.util import warning as c3warn
 from cogent3.util.deserialise import register_deserialiser
@@ -129,14 +130,17 @@ class AlphabetABC(ABC):
     def missing_index(self) -> OptInt: ...
 
     @abstractmethod
-    def to_rich_dict(self) -> dict: ...
+    def to_rich_dict(self, for_pickle: bool = False) -> dict[str, typing.Any]: ...
 
     @abstractmethod
     def to_json(self) -> str: ...
 
     @classmethod
     @abstractmethod
-    def from_rich_dict(cls, data: dict) -> None: ...
+    def from_rich_dict(cls, data: dict) -> typing_extensions.Self: ...
+
+    def __deepcopy__(self, memo):
+        return type(self).from_rich_dict(self.to_rich_dict())
 
     @property
     def moltype(self) -> typing.Union["MolType", None]:
@@ -455,17 +459,20 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         """returns seq as a byte string"""
         return self._arr2bytes(seq)
 
-    def to_rich_dict(self) -> dict:
+    def to_rich_dict(self, for_pickle: bool = False) -> dict[str, typing.Any]:
         """returns a serialisable dictionary"""
         from cogent3._version import __version__
 
-        return {
+        data = {
             "chars": list(self),
             "gap": self.gap_char,
             "missing": self.missing_char,
-            "version": __version__,
-            "type": get_object_provenance(self),
         }
+        if not for_pickle:
+            data["type"] = get_object_provenance(self)
+            data["version"] = __version__
+
+        return data
 
     def to_json(self) -> str:
         """returns a serialisable string"""
@@ -893,19 +900,21 @@ class KmerAlphabet(tuple, AlphabetABC, KmerAlphabetABC):
         max_val = max(self.missing_index or 0, self.gap_index or 0, self.num_canonical)
         return seq.min() >= 0 and seq.max() <= max_val if len(seq) else True
 
-    def to_rich_dict(self) -> dict:
+    def to_rich_dict(self, for_pickle: bool = False) -> dict:
         """returns a serialisable dictionary"""
         from cogent3._version import __version__
 
-        return {
+        data = {
             "words": list(self),
             "monomers": self.monomers.to_rich_dict(),
             "k": self.k,
             "gap": self.gap_char,
             "missing": self.missing_char,
-            "version": __version__,
-            "type": get_object_provenance(self),
         }
+        if not for_pickle:
+            data["type"] = get_object_provenance(self)
+            data["version"] = __version__
+        return data
 
     def to_json(self) -> str:
         """returns a serialisable string"""
@@ -1070,21 +1079,23 @@ class SenseCodonAlphabet(tuple, AlphabetABC):
         words = tuple(self) + (gap_char,)
         return self.__class__(words=words, monomers=monomers, gap=gap_char)
 
-    def to_rich_dict(self):
+    def to_rich_dict(self, for_pickle: bool = False):
         from cogent3._version import __version__
 
-        return {
+        data = {
             "monomers": self.monomers.to_rich_dict(),
-            "type": get_object_provenance(self),
-            "version": __version__,
             "words": list(self),
         }
+        if not for_pickle:
+            data["type"] = get_object_provenance(self)
+            data["version"] = __version__
+        return data
 
     def to_json(self):
         return json.dumps(self.to_rich_dict())
 
     @classmethod
-    def from_rich_dict(cls, data):
+    def from_rich_dict(cls, data: dict[str, typing.Any]):
         data["monomers"] = deserialise_char_alphabet(data["monomers"])
         data.pop("type", None)
         data.pop("version", None)
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index 81627ca42..a55a14554 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -1,6 +1,7 @@
 import dataclasses
 import functools
 import itertools
+import json
 import typing
 from collections import defaultdict
 from string import ascii_letters
@@ -9,6 +10,8 @@ import numpy
 
 from cogent3.core import new_alphabet, new_sequence
 from cogent3.data.molecular_weight import DnaMW, ProteinMW, RnaMW, WeightCalculator
+from cogent3.util.deserialise import register_deserialiser
+from cogent3.util.misc import get_object_provenance
 
 if typing.TYPE_CHECKING:
     from cogent3.core.sequence import SeqViewABC
@@ -1258,6 +1261,20 @@ class MolType:
             seq[index] = f"[{''.join(sorted(expanded))}]"
         return "".join(seq)
 
+    def to_rich_dict(self, **kwargs):
+        data = dict(type=get_object_provenance(self), moltype=self.label)
+        return data
+
+    def to_json(self):
+        """returns result of json formatted string"""
+        data = self.to_rich_dict()
+        return json.dumps(data)
+
+
+@register_deserialiser(get_object_provenance(MolType))
+def deserialise_new_moltype(data, **kwargs):
+    return get_moltype(data.pop("moltype"))
+
 
 def _make_moltype_dict() -> dict[str, MolType]:
     """make a dictionary of local name space molecular types"""


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index aa44cd18c..c995822b4 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -5482,3 +5482,24 @@ def test_slice_preserves_selected_names(DATA_DIR):
     aln = aln.take_seqs(seqnames)
     aln = aln[:1000]
     assert set(aln.names) == set(seqnames)
+
+
+def test_aligned_from_indel_map_and_seqs():
+    dna = new_moltype.get_moltype("dna")
+    seq = dna.make_seq(seq="AC--GTC", name="s1")
+    im, s = seq.parse_out_gaps()
+    al = new_alignment.Aligned.from_map_and_seq(im, s)
+    assert al.name == "s1"
+    assert str(al) == str(seq)
+
+
+def test_aligned_from_indel_map_and_aligned_seq_view():
+    aln = new_alignment.make_aligned_seqs({"s1": "AC--GTC"}, moltype="dna")
+    al = aln.seqs["s1"]
+    new_map = type(al.data.map)(
+        cum_gap_lengths=numpy.array([3, 5], dtype=numpy.int32),
+        gap_pos=numpy.array([0, 2], dtype=numpy.int32),
+        parent_length=4,
+    )
+    new_al = new_alignment.Aligned.from_map_and_aligned_data_view(new_map, al.data)
+    assert str(new_al) == "---AC--GTC"
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index c3eb7ff84..654243c32 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -646,3 +646,16 @@ def test_moltype_coerce_seqs():
     # no coercion in protein seq
     prot = new_moltype.get_moltype("protein")
     assert str(prot.make_seq(seq=rna_seq)) == rna_seq
+
+
+@pytest.mark.parametrize(
+    "moltype", ("protein", "text", "bytes", "protein_with_stop", "dna", "rna")
+)
+def test_json_roundtrip_moltype(moltype):
+    from cogent3.util.deserialise import deserialise_object
+
+    mt = new_moltype.get_moltype(moltype)
+    got = deserialise_object(mt.to_json())
+    # because our moltype instances are singletons,
+    # we expect the same object
+    assert got is mt