From 94b36960815d5cd01a3d163cd6e113007d56c3b5 Mon Sep 17 00:00:00 2001
From: Matt Stone <matt@fulcrumgenomics.com>
Date: Fri, 9 Feb 2024 09:38:09 -0500
Subject: [PATCH] feat: add SupplementaryAlignment constructor from read (#85)

* feat: add SupplementaryAlignment constructor from read

* refactor: use cls

* doc: update docstring
---
 fgpyo/sam/__init__.py                         | 19 +++++++++++++++++
 .../tests/test_supplementary_alignments.py    | 21 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/fgpyo/sam/__init__.py b/fgpyo/sam/__init__.py
index caceadac..bb9fa6a3 100644
--- a/fgpyo/sam/__init__.py
+++ b/fgpyo/sam/__init__.py
@@ -163,6 +163,7 @@
 from typing import Optional
 from typing import Tuple
 from typing import Union
+from typing import cast
 
 import attr
 import pysam
@@ -575,6 +576,24 @@ def parse_sa_tag(tag: str) -> List["SupplementaryAlignment"]:
         """
         return [SupplementaryAlignment.parse(a) for a in tag.split(";") if len(a) > 0]
 
+    @classmethod
+    def from_read(cls, read: pysam.AlignedSegment) -> List["SupplementaryAlignment"]:
+        """
+        Construct a list of SupplementaryAlignments from the SA tag in a pysam.AlignedSegment.
+
+        Args:
+            read: An alignment. The presence of the "SA" tag is not required.
+
+        Returns:
+            A list of all SupplementaryAlignments present in the SA tag.
+            If the SA tag is not present, or it is empty, an empty list will be returned.
+        """
+        if read.has_tag("SA"):
+            sa_tag: str = cast(str, read.get_tag("SA"))
+            return cls.parse_sa_tag(sa_tag)
+        else:
+            return []
+
 
 def isize(r1: AlignedSegment, r2: AlignedSegment) -> int:
     """Computes the insert size for a pair of records."""
diff --git a/fgpyo/sam/tests/test_supplementary_alignments.py b/fgpyo/sam/tests/test_supplementary_alignments.py
index f53efb59..42f4ce6d 100644
--- a/fgpyo/sam/tests/test_supplementary_alignments.py
+++ b/fgpyo/sam/tests/test_supplementary_alignments.py
@@ -2,6 +2,7 @@
 
 from fgpyo.sam import Cigar
 from fgpyo.sam import SupplementaryAlignment
+from fgpyo.sam.builder import SamBuilder
 
 
 def test_supplementary_alignment() -> None:
@@ -48,3 +49,23 @@ def test_format_supplementary_alignment() -> None:
     for sa_string in ["chr1,123,-,100M50S,60,4", "chr1,123,+,100M50S,60,3"]:
         sa = SupplementaryAlignment.parse(sa_string)
         assert str(sa) == sa_string
+
+
+def test_from_read() -> None:
+    """Test that we can construct a SupplementaryAlignment from an AlignedSegment."""
+
+    builder = SamBuilder()
+
+    read = builder.add_single()
+    assert SupplementaryAlignment.from_read(read) == []
+
+    s1 = "chr1,123,+,50S100M,60,0"
+    s2 = "chr2,456,-,75S75M,60,1"
+    sa1 = SupplementaryAlignment("chr1", 122, True, Cigar.from_cigarstring("50S100M"), 60, 0)
+    sa2 = SupplementaryAlignment("chr2", 455, False, Cigar.from_cigarstring("75S75M"), 60, 1)
+
+    read = builder.add_single(attrs={"SA": f"{s1};"})
+    assert SupplementaryAlignment.from_read(read) == [sa1]
+
+    read = builder.add_single(attrs={"SA": f"{s1};{s2};"})
+    assert SupplementaryAlignment.from_read(read) == [sa1, sa2]