From de0a055b9bde46a6ce60dfd51796a16191144cce Mon Sep 17 00:00:00 2001
From: Matt Stone <matt@fulcrumgenomics.com>
Date: Tue, 17 Oct 2023 11:00:38 -0400
Subject: [PATCH] feat: add Template.write_to()

---
 fgpyo/sam/__init__.py                     | 20 +++++++++++++++
 fgpyo/sam/tests/test_template_iterator.py | 31 +++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/fgpyo/sam/__init__.py b/fgpyo/sam/__init__.py
index f9e14776..49e27f2d 100644
--- a/fgpyo/sam/__init__.py
+++ b/fgpyo/sam/__init__.py
@@ -843,6 +843,26 @@ def all_recs(self) -> Iterator[AlignedSegment]:
             for rec in recs:
                 yield rec
 
+    def write_to(
+        self,
+        writer: SamFile,
+        primary_only: bool = False,
+    ) -> None:
+        """Write the records associated with the template to file.
+
+        Args:
+            writer: An open, writable AlignmentFile.
+            primary_only: If True, only write primary alignments.
+        """
+
+        if primary_only:
+            rec_iter = self.primary_recs()
+        else:
+            rec_iter = self.all_recs()
+
+        for rec in rec_iter:
+            writer.write(rec)
+
 
 class TemplateIterator(Iterator[Template]):
     """
diff --git a/fgpyo/sam/tests/test_template_iterator.py b/fgpyo/sam/tests/test_template_iterator.py
index cdefc35a..a28c2dbe 100644
--- a/fgpyo/sam/tests/test_template_iterator.py
+++ b/fgpyo/sam/tests/test_template_iterator.py
@@ -1,4 +1,9 @@
+from pathlib import Path
+
 from fgpyo.sam import Template
+from fgpyo.sam import TemplateIterator
+from fgpyo.sam import reader
+from fgpyo.sam import writer
 from fgpyo.sam.builder import SamBuilder
 
 
@@ -64,3 +69,29 @@ def test_to_templates() -> None:
     assert len(template2.r2_secondaries) == 0
     assert len(list(template2.primary_recs())) == 1
     assert len(list(template2.all_recs())) == 2
+
+
+def test_write_template(tmp_path: Path,) -> None:
+    builder = SamBuilder()
+    template = Template.build([
+        *builder.add_pair(name="r1", chrom="chr1", start1=100, start2=200),
+        builder.add_single(name="r1", chrom="chr1", start=350, supplementary=True),
+    ])
+
+    bam_path = tmp_path / "test.bam"
+
+    # Test writing of all records
+    with writer(bam_path, header=builder._samheader) as bam_writer:
+        template.write_to(bam_writer)
+
+    with reader(bam_path) as bam_reader:
+        template = next(TemplateIterator(bam_reader))
+        assert len([r for r in template.all_recs()]) == 3
+
+    # Test primary-only
+    with writer(bam_path, header=builder._samheader) as bam_writer:
+        template.write_to(bam_writer, primary_only=True)
+
+    with reader(bam_path) as bam_reader:
+        template = next(TemplateIterator(bam_reader))
+        assert len([r for r in template.all_recs()]) == 2