feat: add Template.write_to()

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]):
     """

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,33 @@ 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