htslib 0.4.2 → 0.6.0

data/lib/hts/bcf.rb

@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
 require_relative "../htslib"
-
 require_relative "hts"
+require_relative "native"
 require_relative "bcf/errors"
 require_relative "bcf/header"
 require_relative "bcf/info"
@@ -10,35 +10,36 @@ require_relative "bcf/format"
 require_relative "bcf/record"
 
 module HTS
-  # A class for working with VCF, BCF files.
   class Bcf < Hts
     include Enumerable
 
-    attr_reader :file_name, :index_name, :mode, :header, :nthreads
+    def self.filter_records(records, rid: nil, beg: nil, end_: nil, min_qual: nil, filter_id: nil)
+      Array(records).select do |record|
+        (rid.nil? || record.rid == Integer(rid)) &&
+          (beg.nil? || record.endpos > Integer(beg)) && (end_.nil? || record.pos < Integer(end_)) &&
+          (min_qual.nil? || (!record.qual.nan? && record.qual >= Float(min_qual))) &&
+          (filter_id.nil? || record.filter_id?(filter_id))
+      end
+    end
 
-    def self.open(*args, **kw)
-      file = new(*args, **kw) # do not yield
+    attr_reader :file_name, :index_name, :mode, :header, :nthreads, :unpack
+
+    def self.open(*args, **keywords)
+      file = new(*args, **keywords)
       return file unless block_given?
 
       begin
-        yield file
+        result = yield file
       ensure
         file.close
       end
-      file
+      result
     end
 
     def self.build_index(file_name, index_name = nil, min_shift = 14, threads = 0, verbose = true)
-      if verbose
-        if index_name
-          warn "Create index for #{file_name} to #{index_name}"
-        else
-          warn "Create index for #{file_name}"
-        end
-      end
-
-      case LibHTS.bcf_index_build3(file_name, index_name, min_shift, threads)
-      when 0 # successful
+      warn(index_name ? "Create index for #{file_name} to #{index_name}" : "Create index for #{file_name}") if verbose
+      case Native::BcfFileHandle.build_index(file_name, index_name, min_shift, threads)
+      when 0 then nil
       when -1 then raise IndexError, "Indexing failed for #{file_name}"
       when -2 then raise IndexError, "Opening #{file_name} failed while building the index"
       when -3 then raise IndexError, "#{file_name} is not in an indexable format"
@@ -47,86 +48,117 @@ module HTS
       end
     end
 
-    def initialize(file_name, mode = "r", index: nil, threads: nil,
-                   build_index: false, subset: nil)
-      if block_given?
-        message = "HTS::Bcf.new() does not take block; Please use HTS::Bcf.open() instead"
-        raise message
-      end
-
-      # NOTE: Do not check for the existence of local files, since file_names may be remote URIs.
+    def initialize(file_name, mode = "r", index: nil, threads: nil, build_index: false,
+                   subset: nil, samples: nil, unpack: :all)
+      raise "HTS::Bcf.new() does not take block; Please use HTS::Bcf.open() instead" if block_given?
+      raise ArgumentError, "specify either samples: or subset:, not both" if samples && subset
 
-      @file_name  = file_name
+      subset = samples unless samples.nil?
+      @unpack = unpack.to_sym
+      @max_unpack = resolve_max_unpack(@unpack)
+      @file_name = file_name
       @index_name = index
-      @mode       = mode
-      @nthreads   = threads
-      @hts_file   = LibHTS.hts_open(@file_name, mode)
-
-      raise OpenError, "Failed to open #{@file_name}" if @hts_file.null?
-
+      @mode = mode
+      @nthreads = threads
+      @index_load_attempted = false
+      @native = Native::BcfFileHandle.open(@file_name, mode)
       set_threads(threads) if threads
+      if subset && mode.start_with?("w")
+        raise SubsetError,
+              "Sample subsetting is only available when reading BCF/VCF files"
+      end
+      if writing?
+        @auto_index_on_close = build_index
+        @index_name_on_close = index
+        return
+      end
 
-      raise SubsetError, "Sample subsetting is only available when reading BCF/VCF files" if subset && @mode[0] == "w"
-
-      return if @mode[0] == "w"
-
-      @read_header = Bcf::Header.new(@hts_file)
+      @read_header = Header.new(@native.read_header)
       @header = subset ? @read_header.subset(subset) : @read_header
-      build_index(index) if build_index
-      @idx = load_index(index)
+      if build_index
+        build_index(index)
+        load_index(index)
+      elsif index
+        load_index(index)
+      end
       @start_position = tell
+    rescue Errno::ENOENT
+      raise OpenError, "Failed to open #{@file_name}"
     end
 
     def build_index(index_name = nil, min_shift: 14, verbose: true)
       check_closed
-
       self.class.build_index(@file_name, index_name, min_shift, @nthreads || 0, verbose)
-      self # for method chaining
+      @index_name = index_name
+      @index_load_attempted = false
+      self
     end
 
     def load_index(index_name = nil)
-      check_closed
+      return self if try_load_index(index_name)
 
-      if file_format == "vcf"
-        @index_format = :tabix
-        if index_name
-          LibHTS.tbx_index_load2(@file_name, index_name)
-        else
-          LibHTS.tbx_index_load3(@file_name, nil, 2)
-        end
-      elsif index_name
-        @index_format = :bcf
-        LibHTS.bcf_index_load2(@file_name, index_name)
-      else
-        @index_format = :bcf
-        LibHTS.bcf_index_load3(@file_name, nil, 2)
-      end
+      raise MissingIndexError, "Failed to load index #{index_name || "for #{@file_name}"}"
     end
 
-    def index_loaded?
+    def try_load_index(index_name = nil)
       check_closed
+      @index_name = index_name
+      @index_load_attempted = true
+      @native.load_index(index_name)
+    end
 
-      !@idx.null?
+    def index_loaded?
+      check_closed
+      @native.index_loaded?
     end
 
     def close
-      if @idx && !@idx.null?
-        case @index_format
-        when :bcf
-          LibHTS.hts_idx_destroy(@idx)
-        when :tabix
-          @idx.close
-        end
+      was_closed = closed?
+      result = @native&.close
+      raise WriteError, "Failed to close #{@file_name}: buffered output may be incomplete" if writing? && result&.negative?
+
+      if writing? && @auto_index_on_close && !was_closed
+        @auto_index_on_close = false
+        self.class.build_index(@file_name, @index_name_on_close, 14, @nthreads || 0, false)
       end
-      @idx = nil
-      super
+
+      nil
+    end
+    def closed? = @native.nil? || @native.closed?
+    def file_format = @native.file_format
+    def file_format_version = @native.file_format_version
+    def seek(offset) = @native.seek(offset)
+    def tell = @native.tell
+
+    def rewind
+      raise "Cannot rewind: no start position" unless @start_position
+
+      result = seek(@start_position)
+      raise "Failed to rewind: #{result}" if result.negative?
+
+      tell
+    end
+
+    def set_threads(count = nil)
+      if count.nil?
+        require "etc"
+        count = [Etc.nprocessors - 1, 1].max
+      end
+      raise TypeError unless count.is_a?(Integer)
+      raise ArgumentError, "Number of threads must be positive" if count < 1
+      raise "Failed to set number of threads: #{count}" if @native.set_threads(count).negative?
+
+      @nthreads = count
+      self
     end
 
     def write_header(header)
       check_closed
-
       @header = header.dup
-      LibHTS.bcf_hdr_write(@hts_file, header)
+      result = @native.write_header(header.__send__(:native_handle))
+      raise HeaderError, "Failed to write BCF header" if result.negative?
+
+      nil
     end
 
     def header=(header)
@@ -135,32 +167,27 @@ module HTS
 
     def write(record)
       check_closed
+      result = @native.write(header.__send__(:native_handle), record.__send__(:native_handle))
+      raise "Failed to write record" if result.negative?
 
-      # record = record.dup
-      r = LibHTS.bcf_write(@hts_file, header, record)
-      raise "Failed to write record" if r < 0
+      nil
     end
 
-    def <<(var)
-      write(var)
+    def <<(record)
+      write(record)
+      self
     end
 
     def nsamples
       check_closed
-
       header.nsamples
     end
 
     def samples
       check_closed
-
       header.samples
     end
 
-    # @!macro [attach] define_getter
-    #   @method $1
-    #   Get $1 array
-    #   @return [Array] the $1 array
     define_getter :chrom
     define_getter :pos
     define_getter :endpos
@@ -172,35 +199,28 @@ module HTS
 
     def info(key = nil)
       check_closed
-      position = tell
       raise NotImplementedError unless key
 
-      ary = map { |r| r.info(key) }
-
-      # ary = each_copy.map { |r| r.info }
-      # ary = map { |r| r.info.clone }
-
-      seek(position)
-      ary
+      position = tell
+      map { |record| record.info(key) }.tap { seek(position) if position }
     end
+    alias info_array info
 
     def format(key = nil)
       check_closed
-      position = tell
       raise NotImplementedError unless key
 
-      ary = map { |r| r.format(key) }
+      position = tell
+      map { |record| record.format(key) }.tap { seek(position) if position }
+    end
+    alias format_array format
 
-      # ary = each_copy.map { |r| r.format }
-      # ary = map { |r| r.format.clone }
+    def collect_records = each(copy: true).to_a
 
-      seek(position)
-      ary
-    end
+    # Materialize independent, owning records rather than retaining the
+    # reused Record yielded by the allocation-conscious default iterator.
+    def to_a = collect_records
 
-    # @!macro [attach] define_iterator
-    #   @method each_$1
-    #   Get $1 iterator
     define_iterator :chrom
     define_iterator :pos
     define_iterator :endpos
@@ -211,47 +231,40 @@ module HTS
     define_iterator :filter
 
     def each_info(key)
-      check_closed
       return to_enum(__method__, key) unless block_given?
 
-      each do |r|
-        yield r.info(key)
-      end
+      each { |record| yield record.info(key) }
+      self
     end
 
     def each_format(key)
-      check_closed
       return to_enum(__method__, key) unless block_given?
 
-      each do |r|
-        yield r.format(key)
-      end
+      each { |record| yield record.format(key) }
+      self
     end
 
-    def each(copy: false, &block)
-      if copy
-        each_record_copy(&block)
-      else
-        each_record_reuse(&block)
-      end
-    end
+    def each(copy: false, &block) = copy ? each_record_copy(&block) : each_record_reuse(&block)
 
     def query(region, beg = nil, end_ = nil, copy: false, &block)
       check_closed
-
-      raise MissingIndexError, "Index file is required to call the query method for #{@file_name}" unless index_loaded?
+      unless ensure_index_loaded
+        raise MissingIndexError, "Index file is required to call the query method for #{@file_name}"
+      end
 
       case region
       when Array
         raise ArgumentError, "beg and end must not be specified when region is an Array" unless beg.nil? && end_.nil?
+        return to_enum(__method__, region, copy:) unless block
 
-        query_regions(region, copy:, &block)
+        region.each { |item| query(item, copy:, &block) }
+        self
       else
         if beg && end_
-          tid = header.name2id(region)
-          queryi(tid, beg, end_, copy:, &block)
+          iterate_query(@native.query_interval(read_header_native, header.name2id(region), beg, end_), copy, region,
+                        &block)
         elsif beg.nil? && end_.nil?
-          querys(region, copy:, &block)
+          iterate_query(@native.query_region(read_header_native, region), copy, region, &block)
         else
           raise ArgumentError, "beg and end must be specified together"
         end
@@ -260,241 +273,85 @@ module HTS
 
     private
 
-    def queryi(tid, beg, end_, copy: false, &block)
-      if copy
-        queryi_copy(tid, beg, end_, &block)
-      else
-        queryi_reuse(tid, beg, end_, &block)
-      end
-    end
-
-    def querys(region, copy: false, &block)
-      if copy
-        querys_copy(region, &block)
-      else
-        querys_reuse(region, &block)
-      end
-    end
-
-    def query_regions(regions, copy: false, &block)
-      if copy
-        query_regions_copy(regions, &block)
-      else
-        query_regions_reuse(regions, &block)
-      end
-    end
-
-    def queryi_reuse(tid, beg, end_, &block)
-      return to_enum(__method__, tid, beg, end_) unless block_given?
-
-      return queryi_reuse_vcf(tid, beg, end_, &block) if tabix_index?
-
-      qiter = LibHTS.bcf_itr_queryi(@idx, tid, beg, end_)
-      raise QueryError, "Failed to query region #{tid}:#{beg}-#{end_} in #{@file_name}" if qiter.null?
-
-      query_reuse_yield(qiter, &block)
-      self
-    end
-
-    def querys_reuse(region, &block)
-      return to_enum(__method__, region) unless block_given?
-
-      return querys_reuse_vcf(region, &block) if tabix_index?
-
-      qiter = LibHTS.bcf_itr_querys(@idx, read_header, region)
-      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" if qiter.null?
-
-      query_reuse_yield(qiter, &block)
-      self
-    end
-
-    def query_regions_reuse(regions, &block)
-      return to_enum(__method__, regions) unless block_given?
-
-      regions.each do |region|
-        querys_reuse(region, &block)
-      end
-      self
-    end
-
-    def query_reuse_yield(qiter)
-      bcf1 = LibHTS.bcf_init
-      record = Record.new(header, bcf1)
-      begin
-        loop do
-          slen = LibHTS.bcf_itr_next(@hts_file, qiter, bcf1)
-          break if slen == -1
-          raise if slen < -1
-
-          apply_subset!(record)
-          yield record
-        end
-      ensure
-        LibHTS.bcf_itr_destroy(qiter)
-      end
-    end
-
-    def queryi_copy(tid, beg, end_, &block)
-      return to_enum(__method__, tid, beg, end_) unless block_given?
-
-      return queryi_copy_vcf(tid, beg, end_, &block) if tabix_index?
-
-      qiter = LibHTS.bcf_itr_queryi(@idx, tid, beg, end_)
-      raise QueryError, "Failed to query region #{tid}:#{beg}-#{end_} in #{@file_name}" if qiter.null?
-
-      query_copy_yield(qiter, &block)
-      self
-    end
-
-    def querys_copy(region, &block)
-      return to_enum(__method__, region) unless block_given?
+    def writing? = @mode&.start_with?("w", "a")
 
-      return querys_copy_vcf(region, &block) if tabix_index?
+    def ensure_index_loaded
+      return true if index_loaded?
+      return false if @index_load_attempted
 
-      qiter = LibHTS.bcf_itr_querys(@idx, read_header, region)
-      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" if qiter.null?
-
-      query_copy_yield(qiter, &block)
-      self
+      load_index(@index_name)
     end
 
-    def query_regions_copy(regions, &block)
-      return to_enum(__method__, regions) unless block_given?
+    def native_handle = @native
+    def read_header_native = (@read_header || @header).__send__(:native_handle)
 
-      regions.each do |region|
-        querys_copy(region, &block)
-      end
-      self
-    end
+    def each_record_reuse
+      return to_enum(__method__) unless block_given?
 
-    def query_copy_yield(qiter)
-      bcf1 = LibHTS.bcf_init
-      record = Record.new(header, bcf1)
+      record = Record.new(header)
+      prepare_record(record)
       loop do
-        slen = LibHTS.bcf_itr_next(@hts_file, qiter, bcf1)
-        break if slen == -1
-        raise if slen < -1
+        result = @native.read(read_header_native, record.__send__(:native_handle))
+        break if result == -1
+        raise QueryError, "Failed to read variant record from #{@file_name}" if result < -1
 
         apply_subset!(record)
-        yield record.dup
+        yield record
       end
-    ensure
-      LibHTS.bcf_itr_destroy(qiter)
-    end
-
-    def tabix_index?
-      @index_format == :tabix
-    end
-
-    def queryi_reuse_vcf(tid, beg, end_, &block)
-      qiter = LibHTS.tbx_itr_queryi(@idx, tid, beg, end_)
-      raise QueryError, "Failed to query region #{tid}:#{beg}-#{end_} in #{@file_name}" if qiter.null?
-
-      query_reuse_yield_vcf(qiter, &block)
-      self
-    end
-
-    def querys_reuse_vcf(region, &block)
-      qiter = LibHTS.tbx_itr_querys(@idx, region)
-      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" if qiter.null?
-
-      query_reuse_yield_vcf(qiter, &block)
       self
     end
 
-    def query_reuse_yield_vcf(qiter)
-      line = LibHTS::KString.new
-      bcf1 = LibHTS.bcf_init
-      record = Record.new(header, bcf1)
-      begin
-        while (slen = LibHTS.tbx_itr_next(@hts_file, @idx, qiter, line)) >= 0
-          raise QueryError, "Failed to parse VCF record in #{@file_name}" if LibHTS.vcf_parse(line, read_header,
-                                                                                              bcf1) < 0
-
-          apply_subset!(record)
-          yield record
-        end
-        raise if slen < -1
-      ensure
-        line.free_buffer
-        LibHTS.hts_itr_destroy(qiter)
-      end
-    end
-
-    def queryi_copy_vcf(tid, beg, end_, &block)
-      qiter = LibHTS.tbx_itr_queryi(@idx, tid, beg, end_)
-      raise QueryError, "Failed to query region #{tid}:#{beg}-#{end_} in #{@file_name}" if qiter.null?
+    def each_record_copy
+      return to_enum(__method__) unless block_given?
 
-      query_copy_yield_vcf(qiter, &block)
+      each_record_reuse { |record| yield record.dup }
       self
     end
 
-    def querys_copy_vcf(region, &block)
-      qiter = LibHTS.tbx_itr_querys(@idx, region)
-      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" if qiter.null?
+    def iterate_query(iterator, copy, region)
+      return to_enum(__method__, iterator, copy, region) unless block_given?
+      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" unless iterator
 
-      query_copy_yield_vcf(qiter, &block)
-      self
-    end
-
-    def query_copy_yield_vcf(qiter)
-      line = LibHTS::KString.new
+      record = Record.new(header)
+      prepare_record(record)
       begin
-        while (slen = LibHTS.tbx_itr_next(@hts_file, @idx, qiter, line)) >= 0
-          bcf1 = LibHTS.bcf_init
-          raise QueryError, "Failed to parse VCF record in #{@file_name}" if LibHTS.vcf_parse(line, read_header,
-                                                                                              bcf1) < 0
+        loop do
+          result = iterator.next(record.__send__(:native_handle))
+          break if result == -1
+          raise QueryError, "Failed to parse/query record in #{@file_name}" if result < -1
 
-          record = Record.new(header, bcf1)
           apply_subset!(record)
-          yield record
+          yield(copy ? record.dup : record)
         end
-        raise if slen < -1
       ensure
-        line.free_buffer
-        LibHTS.hts_itr_destroy(qiter)
-      end
-    end
-
-    def each_record_reuse
-      check_closed
-
-      return to_enum(__method__) unless block_given?
-
-      bcf1 = LibHTS.bcf_init
-      record = Record.new(header, bcf1)
-      while LibHTS.bcf_read(@hts_file, read_header, bcf1) != -1
-        apply_subset!(record)
-        yield record
+        iterator.close
       end
       self
     end
 
-    def each_record_copy
-      check_closed
+    def apply_subset!(record)
+      return unless header&.subset?
 
-      return to_enum(__method__) unless block_given?
+      map = header.__send__(:subset_imap)
+      result = record.__send__(:native_handle).subset(header.__send__(:native_handle), map)
+      return unless result.negative?
 
-      bcf1 = LibHTS.bcf_init
-      record = Record.new(header, bcf1)
-      while LibHTS.bcf_read(@hts_file, read_header, bcf1) != -1
-        apply_subset!(record)
-        yield record.dup
-      end
-      self
+      raise SubsetError,
+            "Failed to subset samples #{header.subset_samples.inspect} while reading #{@file_name}"
     end
 
-    def read_header
-      @read_header || header
+    def resolve_max_unpack(level)
+      case level
+      when :all, :format then 15
+      when :site_only, :info then 4
+      when :filter then 2
+      when :string, :alleles then 1
+      else raise ArgumentError, "unknown unpack level: #{level.inspect}"
+      end
     end
 
-    def apply_subset!(record)
-      return unless header.subset?
-
-      rc = LibHTS.bcf_subset(header.struct, record.struct, header.subset_sample_count, header.subset_imap_pointer || ::FFI::Pointer::NULL)
-      return if rc >= 0
-
-      raise SubsetError, "Failed to subset samples #{header.subset_samples.inspect} while reading #{@file_name}"
+    def prepare_record(record)
+      record.__send__(:native_handle).max_unpack = @max_unpack
     end
   end
 end