Enhance pileup and mpileup iterators with RAII-style resource management

Author: kojix2

lib/hts/bam.rb

@@ -228,20 +228,24 @@ def query(region, beg = nil, end_ = nil, copy: false, &block)
     end
 
     # Pileup iterator over this file. Optional region can be specified.
-    # When a block is given, yields PileupColumn and returns self.
-    # Without a block, returns an Enumerator.
+    # When a block is given, uses RAII-style and ensures the iterator is closed at block end.
+    # Without a block, returns an Enumerator over a live Pileup instance; caller should close when done.
     #
     # @param region [String, nil] region string like "chr1:100-200"
     # @param beg [Integer, nil]
     # @param end_ [Integer, nil]
     # @param maxcnt [Integer, nil] cap on depth per position
     def pileup(region = nil, beg = nil, end_: nil, maxcnt: nil, &block)
       check_closed
-      piter = Pileup.new(self, region:, beg:, end_: end_, maxcnt: maxcnt)
-      return piter.to_enum(:each) unless block_given?
-
-      piter.each(&block)
-      self
+      if block_given?
+        Pileup.open(self, region:, beg:, end_: end_, maxcnt: maxcnt) do |piter|
+          piter.each(&block)
+        end
+        self
+      else
+        piter = Pileup.new(self, region:, beg:, end_: end_, maxcnt: maxcnt)
+        piter.to_enum(:each)
+      end
     end
 
     private

lib/hts/bam/mpileup.rb

@@ -6,6 +6,22 @@ class Bam < Hts
     class Mpileup
       include Enumerable
 
+      # Usage:
+      #   HTS::Bam::Mpileup.open([bam1, bam2], region: "chr1:1-100") do |mpl|
+      #     mpl.each { |cols| ... }
+      #   end
+      def self.open(*args, **kw)
+        m = new(*args, **kw)
+        return m unless block_given?
+
+        begin
+          yield m
+        ensure
+          m.close
+        end
+        m
+      end
+
       # Normalize inputs to HTS::Bam instances
       # Accepts array of HTS::Bam or filenames (String)
       def initialize(inputs, region: nil, beg: nil, end_: nil, maxcnt: nil, overlaps: false)
@@ -107,39 +123,35 @@ def each
         plp1_size = HTS::LibHTS::BamPileup1.size
         headers   = @bams.map(&:header)
 
-        begin
-          while HTS::LibHTS.bam_mplp64_auto(@iter, tid_ptr, pos_ptr, n_ptr, plp_ptr) > 0
-            tid = tid_ptr.read_int
-            pos = pos_ptr.read_long_long
-
-            counts = n_ptr.read_array_of_int(n)
-            plp_arr = plp_ptr.read_array_of_pointer(n)
-
-            cols = Array.new(n)
-            i = 0
-            while i < n
-              c = counts[i]
-              if c <= 0 || plp_arr[i].null?
-                cols[i] = HTS::Bam::Pileup::PileupColumn.new(tid: tid, pos: pos, alignments: [])
-              else
-                base_ptr = plp_arr[i]
-                aligns = Array.new(c)
-                j = 0
-                while j < c
-                  e_ptr = base_ptr + (j * plp1_size)
-                  entry = HTS::LibHTS::BamPileup1.new(e_ptr)
-                  aligns[j] = HTS::Bam::Pileup::PileupRecord.new(entry, headers[i])
-                  j += 1
-                end
-                cols[i] = HTS::Bam::Pileup::PileupColumn.new(tid: tid, pos: pos, alignments: aligns)
+        while HTS::LibHTS.bam_mplp64_auto(@iter, tid_ptr, pos_ptr, n_ptr, plp_ptr) > 0
+          tid = tid_ptr.read_int
+          pos = pos_ptr.read_long_long
+
+          counts = n_ptr.read_array_of_int(n)
+          plp_arr = plp_ptr.read_array_of_pointer(n)
+
+          cols = Array.new(n)
+          i = 0
+          while i < n
+            c = counts[i]
+            if c <= 0 || plp_arr[i].null?
+              cols[i] = HTS::Bam::Pileup::PileupColumn.new(tid: tid, pos: pos, alignments: [])
+            else
+              base_ptr = plp_arr[i]
+              aligns = Array.new(c)
+              j = 0
+              while j < c
+                e_ptr = base_ptr + (j * plp1_size)
+                entry = HTS::LibHTS::BamPileup1.new(e_ptr)
+                aligns[j] = HTS::Bam::Pileup::PileupRecord.new(entry, headers[i])
+                j += 1
               end
-              i += 1
+              cols[i] = HTS::Bam::Pileup::PileupColumn.new(tid: tid, pos: pos, alignments: aligns)
             end
-
-            yield cols
+            i += 1
           end
-        ensure
-          close
+
+          yield cols
         end
 
         self

lib/hts/bam/pileup.rb

@@ -6,6 +6,22 @@ class Bam < Hts
     class Pileup
       include Enumerable
 
+      # Usage:
+      #   HTS::Bam::Pileup.open(bam, region: "chr1:1-100") do |pl|
+      #     pl.each { |col| ... }
+      #   end
+      def self.open(*args, **kw)
+        pu = new(*args, **kw)
+        return pu unless block_given?
+
+        begin
+          yield pu
+        ensure
+          pu.close
+        end
+        pu
+      end
+
       # A column at a reference position with pileup alignments
       PileupColumn = Struct.new(:tid, :pos, :alignments, keyword_init: true) do
         def depth
@@ -135,40 +151,36 @@ def each
         plp1_size    = HTS::LibHTS::BamPileup1.size
         header_local = @header
 
-        begin
-          loop do
-            base_ptr = HTS::LibHTS.bam_plp64_auto(@plp, tid_ptr, pos_ptr, n_ptr)
+        loop do
+          base_ptr = HTS::LibHTS.bam_plp64_auto(@plp, tid_ptr, pos_ptr, n_ptr)
 
-            # When base_ptr is NULL, check n to distinguish EOF (n == 0) from error (n < 0)
-            if base_ptr.null?
-              n = n_ptr.read_int
-              raise "HTSlib pileup error (bam_plp64_auto)" if n < 0
+          # When base_ptr is NULL, check n to distinguish EOF (n == 0) from error (n < 0)
+          if base_ptr.null?
+            n = n_ptr.read_int
+            raise "HTSlib pileup error (bam_plp64_auto)" if n < 0
 
-              break
-            end
+            break
+          end
 
-            tid = tid_ptr.read_int
-            pos = pos_ptr.read_long_long
-            n   = n_ptr.read_int
-
-            # Construct alignment entries with minimal allocations
-            if n.zero?
-              alignments = []
-            else
-              alignments = Array.new(n)
-              i = 0
-              while i < n
-                e_ptr = base_ptr + (i * plp1_size)
-                entry = HTS::LibHTS::BamPileup1.new(e_ptr)
-                alignments[i] = PileupRecord.new(entry, header_local)
-                i += 1
-              end
+          tid = tid_ptr.read_int
+          pos = pos_ptr.read_long_long
+          n   = n_ptr.read_int
+
+          # Construct alignment entries with minimal allocations
+          if n.zero?
+            alignments = []
+          else
+            alignments = Array.new(n)
+            i = 0
+            while i < n
+              e_ptr = base_ptr + (i * plp1_size)
+              entry = HTS::LibHTS::BamPileup1.new(e_ptr)
+              alignments[i] = PileupRecord.new(entry, header_local)
+              i += 1
             end
-
-            yield PileupColumn.new(tid: tid, pos: pos, alignments: alignments)
           end
-        ensure
-          close
+
+          yield PileupColumn.new(tid: tid, pos: pos, alignments: alignments)
         end
 
         self