From b0b4227de8acff8a47f48d2afb06190c1b9850ad Mon Sep 17 00:00:00 2001
From: Clint Valentine <valentine.clint@gmail.com>
Date: Mon, 6 Jan 2025 20:34:13 -0500
Subject: [PATCH] Add a method for progress logging multiple alignments (#199)

---
 fgpyo/util/logging.py            | 18 ++++++++++++++
 tests/fgpyo/util/test_logging.py | 40 ++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/fgpyo/util/logging.py b/fgpyo/util/logging.py
index d956c716..6603b9a7 100644
--- a/fgpyo/util/logging.py
+++ b/fgpyo/util/logging.py
@@ -38,6 +38,7 @@
 from threading import RLock
 from typing import Any
 from typing import Callable
+from typing import Iterable
 from typing import Literal
 from typing import Optional
 from typing import Union
@@ -163,6 +164,23 @@ def record_alignment(
         else:
             return self.record(rec.reference_name, rec.reference_start + 1)
 
+    def record_alignments(
+        self,
+        recs: Iterable[AlignedSegment],
+    ) -> bool:
+        """Correctly record multiple pysam.AlignedSegments (zero-based coordinates).
+
+        Args:
+            recs: pysam.AlignedSegment objects
+
+        Returns:
+            true if a message was logged, false otherwise
+        """
+        logged_message: bool = False
+        for rec in recs:
+            logged_message = self.record_alignment(rec) or logged_message
+        return logged_message
+
     def _log(
         self,
         refname: Optional[str] = None,
diff --git a/tests/fgpyo/util/test_logging.py b/tests/fgpyo/util/test_logging.py
index 21a6e283..ca8584c1 100644
--- a/tests/fgpyo/util/test_logging.py
+++ b/tests/fgpyo/util/test_logging.py
@@ -4,6 +4,7 @@
 import pytest
 
 from fgpyo import sam
+from fgpyo.sam import Template
 from fgpyo.sam.builder import SamBuilder
 from fgpyo.util.logging import ProgressLogger
 
@@ -59,3 +60,42 @@ def test_record_alignment_mapped_record(record: pysam.AlignedSegment) -> None:
 
     # Assert record is logged
     assert progress.record_alignment(rec=record) is True
+
+
+def test_record_multiple_alignments() -> None:
+    builder: SamBuilder = SamBuilder()
+    (r1, r2) = builder.add_pair(name="x", chrom="chr1", start1=1, start2=2)
+    (r1_secondary, r2_secondary) = builder.add_pair(name="x", chrom="chr1", start1=10, start2=12)
+    r1_secondary.is_secondary = True
+    r2_secondary.is_secondary = True
+    (r1_supplementary, r2_supplementary) = builder.add_pair(
+        name="x", chrom="chr1", start1=4, start2=6
+    )
+    r1_supplementary.is_supplementary = True
+    r2_supplementary.is_supplementary = True
+
+    template = Template.build(builder.to_unsorted_list())
+    expected = Template(
+        name="x",
+        r1=r1,
+        r2=r2,
+        r1_secondaries=[r1_secondary],
+        r2_secondaries=[r2_secondary],
+        r1_supplementals=[r1_supplementary],
+        r2_supplementals=[r2_supplementary],
+    )
+
+    assert template == expected
+
+    # Define instance of ProgressLogger
+    actual: list[str] = []
+
+    progress = ProgressLogger(
+        printer=lambda rec: actual.append(rec), noun="record(s)", verb="recorded", unit=1
+    )
+
+    # Assert record is logged
+    assert progress.record_alignments(recs=template.all_recs()) is True
+
+    # Assert every record was logged
+    assert len(actual) == 6