Repository: cogent3/cogent3
Pull request #1822 (2024-04-24 06:29:12)

ENH: seqcoll.to_fasta supports block_size
Addresses #1795 


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index e72883d40..1bba1b448 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -984,9 +984,20 @@ class _SequenceCollectionBase:
         """returns json formatted string"""
         return json.dumps(self.to_rich_dict())
 
-    def to_fasta(self):
-        """Return alignment in Fasta format"""
-        return alignment_to_fasta(self.to_dict())
+    def to_fasta(self, block_size: int = 60) -> str:
+        """Return alignment in Fasta format.
+
+        Parameters
+        ----------
+        block_size
+            the sequence length to write to each line,
+            by default 60
+
+        Returns
+        -------
+        The Fasta formatted alignment.
+        """
+        return alignment_to_fasta(self.to_dict(), block_size=block_size)
 
     def to_nexus(self, seq_type, wrap=50):
         """
@@ -1192,10 +1203,15 @@ class _SequenceCollectionBase:
         """Return a sequence object for the specified seqname."""
         return self.named_seqs[seqname]
 
-    def to_dict(self):
-        """Returns the alignment as dict of names -> strings.
+    def to_dict(self) -> dict[str, str]:
+        """Returns the alignment as a dict of sequence names -> strings.
+
+        Note the mapping goes to strings, not Sequence objects.
 
-        Note: returns strings, NOT Sequence objects.
+        Returns
+        -------
+        a dict mapping sequence names to a string representation of
+        their sequences.
         """
         align_dict = {}
 
diff --git a/src/cogent3/format/fasta.py b/src/cogent3/format/fasta.py
index 0a4ebf216..28a3bf375 100644
--- a/src/cogent3/format/fasta.py
+++ b/src/cogent3/format/fasta.py
@@ -2,17 +2,45 @@
 """Writer for FASTA sequence format
 """
 
+from typing import Optional
+
 from cogent3.format.util import _AlignmentFormatter
 
 
-def alignment_to_fasta(alignment_dict, block_size=60, order=None):
-    """Returns a Fasta string given an alignment."""
+def alignment_to_fasta(
+    alignment_dict: dict[str, str],
+    block_size: int = 60,
+    order: Optional[list[str]] = None,
+) -> str:
+    """Returns a Fasta string given an alignment.
+
+    Parameters
+        ----------
+        alignment_dict
+            dict of seq_name
+        block_size
+            the sequence length to write to each line,
+            by default 60
+        order
+            optional list of sequence names, which order to print in.
+            Assumes complete and correct list of names,
+            by default None
+
+        Returns
+        -------
+        The alignment in the Fasta format.
+    """
     order = order or []
     return FastaFormatter().format(alignment_dict, block_size, order)
 
 
 class FastaFormatter(_AlignmentFormatter):
-    def format(self, alignment_dict, block_size, order):
+    def format(
+        self,
+        alignment_dict: dict[str, str],
+        block_size: int = 60,
+        order: Optional[list[str]] = None,
+    ) -> str:
         """Format the alignment to Fasta.
 
         Parameters
@@ -21,12 +49,15 @@ class FastaFormatter(_AlignmentFormatter):
             dict of seq_name
         block_size
             the sequence length to write to each line,
-            default is 60
+            by default 60
         order
-            optional list of sequence names, which order to
-            print in.
+            optional list of sequence names, which order to print in.
+            Assumes complete and correct list of names,
+            by default None
 
-        (Assumes complete and correct list of names)
+        Returns
+        -------
+        The alignment in the Fasta format.
         """
         # setup
         if not order:


Test patch:
diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
index 78070ee2b..dcb2ec9cd 100644
--- a/tests/test_core/test_alignment.py
+++ b/tests/test_core/test_alignment.py
@@ -413,8 +413,18 @@ class SequenceCollectionBaseTests(object):
 
     def test_to_fasta(self):
         """SequenceCollection should return correct FASTA string"""
-        aln = self.Class(["AAA", "CCC"])
-        self.assertEqual(aln.to_fasta(), ">seq_0\nAAA\n>seq_1\nCCC\n")
+        aln1 = self.Class(["AAA", "CCC"])
+        self.assertEqual(aln1.to_fasta(), ">seq_0\nAAA\n>seq_1\nCCC\n")
+        self.assertEqual(aln1.to_fasta(block_size=2), ">seq_0\nAA\nA\n>seq_1\nCC\nC\n")
+
+        aln2 = self.Class(["GCATGCAT", "TCAGACGT"])
+        self.assertEqual(aln2.to_fasta(), ">seq_0\nGCATGCAT\n>seq_1\nTCAGACGT\n")
+        self.assertEqual(
+            aln2.to_fasta(block_size=4), ">seq_0\nGCAT\nGCAT\n>seq_1\nTCAG\nACGT\n"
+        )
+        self.assertEqual(
+            aln2.to_fasta(block_size=3), ">seq_0\nGCA\nTGC\nAT\n>seq_1\nTCA\nGAC\nGT\n"
+        )
 
     def test_to_nexus(self):
         """SequenceCollection should return correct Nexus string format"""