csv-utils 0.3.25 → 0.6.0

data/lib/csv_utils/csv_sort.rb

@@ -1,116 +1,156 @@
-require 'fileutils'
-
-# Utility class for sorting the rows for a csv file
-class CSVUtils::CSVSort
-  attr_reader :csv_file,
-              :new_csv_file,
-              :has_headers,
-              :csv_options,
-              :headers
-
-  def initialize(csv_file, new_csv_file, has_headers = true, csv_options = {})
-    @csv_file = csv_file
-    @new_csv_file = new_csv_file
-    @has_headers = has_headers
-    @csv_options = csv_options
-    @csv_part_files = []
-    @files_to_delete = []
-  end
-
-  def sort(batch_size = 100_000, &block)
-    create_sorted_csv_part_files(batch_size, &block)
-    merge_csv_part_files(&block)
-  end
-
-  private
+# frozen_string_literal: true
 
-  def merge_sort_csv_files(src_csv_file1, src_csv_file2, dest_csv_file)
-    src1 = CSV.open(src_csv_file1, 'rb', **csv_options)
-    src2 = CSV.open(src_csv_file2, 'rb', **csv_options)
-    dest = CSV.open(dest_csv_file, 'wb', **csv_options)
-
-    if @headers
-      dest << @headers
-      src1.shift
-      src2.shift
-    end
-
-    row1 = src1.shift
-    row2 = src2.shift
+require 'fileutils'
 
-    append_row1_proc = Proc.new do
-      dest << row1
-      row1 = src1.shift
+module CSVUtils
+  # Sorts a CSV file too large to load into memory with an external merge sort: batches of rows are
+  # sorted into temporary part files next to the new file, which are then merged in pairs.
+  #
+  # @example
+  #   CSVUtils::CSVSort.new('input.csv', 'sorted.csv').sort { |a, b| a[0].to_i <=> b[0].to_i }
+  class CSVSort
+    # @return [String] path of the file to sort
+    attr_reader :csv_file
+    # @return [String] path of the sorted file to write
+    attr_reader :new_csv_file
+    # @return [Boolean] whether the first row is a header row, kept at the top
+    attr_reader :has_headers
+    # @return [Hash] options passed to CSV.open
+    attr_reader :csv_options
+    # @return [Array<String>, nil] the header row, once {#sort} has read it
+    attr_reader :headers
+
+    # @param csv_file [String] path of the file to sort
+    # @param new_csv_file [String] path of the sorted file to write
+    # @param has_headers [Boolean] whether the first row is a header row
+    # @param csv_options [Hash] options passed to CSV.open
+    def initialize(csv_file, new_csv_file, has_headers = true, csv_options = {})
+      @csv_file = csv_file
+      @new_csv_file = new_csv_file
+      @has_headers = has_headers
+      @csv_options = csv_options
+      @csv_part_files = []
     end
 
-    append_row2_proc = Proc.new do
-      dest << row2
-      row2 = src2.shift
+    # Writes the sorted file. Temporary files are removed even when sorting fails.
+    # @param batch_size [Integer] rows held in memory and sorted at a time
+    # @yieldparam row1 [Array<String>]
+    # @yieldparam row2 [Array<String>]
+    # @yieldreturn [Integer] negative, 0 or positive like <=>; without a block rows are compared as arrays
+    # @return [void]
+    def sort(batch_size = 100_000, &block)
+      block ||= proc { |row1, row2| row1 <=> row2 }
+      @csv_part_files = []
+      create_sorted_csv_part_files(batch_size, &block)
+      merge_csv_part_files(&block)
+    ensure
+      delete_csv_part_files
     end
 
-    while row1 || row2
-      if row1.nil?
-        append_row2_proc.call
-      elsif row2.nil?
-        append_row1_proc.call
-      elsif yield(row1, row2) <= 0
-        append_row1_proc.call
-      else
-        append_row2_proc.call
+    private
+
+    # rubocop:disable-next Metrics/MethodLength
+    def merge_sort_csv_files(src_csv_file1, src_csv_file2, dest_csv_file)
+      src1 = CSV.open(src_csv_file1, 'rb', **csv_options)
+      begin
+        src2 = CSV.open(src_csv_file2, 'rb', **csv_options)
+        begin
+          dest = CSV.open(dest_csv_file, 'wb', **csv_options)
+          begin
+            if @headers
+              dest << @headers
+              src1.shift
+              src2.shift
+            end
+
+            row1 = src1.shift
+            row2 = src2.shift
+
+            append_row1_proc = proc do
+              dest << row1
+              row1 = src1.shift
+            end
+
+            append_row2_proc = proc do
+              dest << row2
+              row2 = src2.shift
+            end
+
+            while row1 || row2
+              if row1.nil?
+                append_row2_proc.call
+              elsif row2.nil?
+                append_row1_proc.call
+              elsif yield(row1, row2) <= 0
+                append_row1_proc.call
+              else
+                append_row2_proc.call
+              end
+            end
+          ensure
+            dest.close
+          end
+        ensure
+          src2.close
+        end
+      ensure
+        src1.close
       end
     end
 
-    src1.close
-    src2.close
-    dest.close
-  end
-
-  def create_sorted_csv_part_files(batch_size, &block)
-    src = CSV.open(csv_file, 'rb', **csv_options)
-
-    @headers = src.shift if has_headers
-
-    batch = []
-    create_batch_part_proc = Proc.new do
-      batch.sort!(&block)
-      @csv_part_files << "#{new_csv_file}.part.#{@csv_part_files.size}"
-      CSV.open(@csv_part_files.last, 'wb', **csv_options) do |csv|
-        csv << @headers if @headers
-        batch.each { |row| csv << row }
+    def create_sorted_csv_part_files(batch_size, &block)
+      src = CSV.open(csv_file, 'rb', **csv_options)
+      begin
+        @headers = src.shift if has_headers
+
+        batch = []
+        create_batch_part_proc = proc do
+          batch.sort!(&block)
+          @csv_part_files << "#{new_csv_file}.part.#{@csv_part_files.size}"
+          CSV.open(@csv_part_files.last, 'wb', **csv_options) do |csv|
+            csv << @headers if @headers
+            batch.each { |row| csv << row }
+          end
+          batch = []
+        end
+
+        while (row = src.shift)
+          batch << row
+          create_batch_part_proc.call if batch.size >= batch_size
+        end
+
+        create_batch_part_proc.call if batch.size.positive?
+      ensure
+        src.close
       end
-      batch = []
     end
 
-    while (row = src.shift)
-      batch << row
-      create_batch_part_proc.call if batch.size >= batch_size
-    end
+    def merge_csv_part_files(&)
+      file_merge_cnt = 0
 
-    create_batch_part_proc.call if batch.size > 0
+      while @csv_part_files.size > 1
+        file_merge_cnt += 1
 
-    src.close
-  end
+        # inputs stay in the list until merged, so a failed merge still cleans them up
+        csv_part_file1, csv_part_file2 = @csv_part_files.first(2)
+        @csv_part_files << "#{new_csv_file}.merge.#{file_merge_cnt}"
 
-  def merge_csv_part_files(&block)
-    file_merge_cnt = 0
+        merge_sort_csv_files(csv_part_file1, csv_part_file2, @csv_part_files.last, &)
 
-    while @csv_part_files.size > 1
-      file_merge_cnt += 1
-
-      csv_part_file1 = @csv_part_files.shift
-      csv_part_file2 = @csv_part_files.shift
-      @csv_part_files << "#{new_csv_file}.merge.#{file_merge_cnt}"
-
-      merge_sort_csv_files(csv_part_file1, csv_part_file2, @csv_part_files.last, &block)
+        @csv_part_files.shift(2).each { |file| File.unlink(file) }
+      end
 
-      File.unlink(csv_part_file1)
-      File.unlink(csv_part_file2)
+      if @csv_part_files.size.positive?
+        FileUtils.mv(@csv_part_files.pop, new_csv_file)
+      else
+        FileUtils.cp(@csv_file, new_csv_file)
+      end
     end
 
-    if @csv_part_files.size > 0
-      FileUtils.mv(@csv_part_files.last, new_csv_file)
-    else
-      FileUtils.cp(@csv_file, new_csv_file)
+    # removes the temporary files left behind when sorting fails part way through
+    def delete_csv_part_files
+      @csv_part_files.each { |file| FileUtils.rm_f(file) }
+      @csv_part_files = []
     end
   end
 end

data/lib/csv_utils/csv_transformer.rb

@@ -1,119 +1,166 @@
-# Transforms a CSV given a series of steps
-class CSVUtils::CSVTransformer
-  attr_reader :headers
-
-  def initialize(src_csv, dest_csv, csv_options = {})
-    @src_csv = CSVUtils::CSVWrapper.new(src_csv, 'rb', csv_options)
-    @dest_csv = CSVUtils::CSVWrapper.new(dest_csv, 'wb', csv_options)
-  end
+# frozen_string_literal: true
+
+module CSVUtils
+  # Copies a CSV file through a pipeline of steps that filter, change and add to its rows in batches.
+  # Each step method returns the transformer so they can be chained, and {#process} runs them.
+  #
+  # Step blocks receive the row, the headers as they were when the step was added, and the value
+  # computed by the last {#additional_data} step for the current batch.
+  #
+  # @example
+  #   CSVUtils::CSVTransformer.new('input.csv', 'output.csv')
+  #     .read_headers
+  #     .select { |row, _headers, _data| row[0].to_i > 100 }
+  #     .append(['Email']) { |row, _headers, _data| ["#{row[1].downcase}@example.com"] }
+  #     .process
+  class CSVTransformer
+    # @return [Array<String>, nil] the headers the output will have after the steps added so far
+    attr_reader :headers
+
+    # @param src_csv [String, CSV, #shift] path or CSV to read; paths are opened and closed by the transformer
+    # @param dest_csv [String, CSV, #<<] path or CSV to write
+    # @param csv_options [Hash] options passed to CSV.open for paths
+    def initialize(src_csv, dest_csv, csv_options = {})
+      @src_csv = CSVUtils::CSVWrapper.new(src_csv, 'rb', csv_options)
+      @dest_csv = CSVUtils::CSVWrapper.new(dest_csv, 'wb', csv_options)
+    end
 
-  def read_headers
-    @headers = @src_csv.shift
-    self
-  end
+    # Reads the first row as the headers. Without it, every row is data and the output has no header row.
+    # @return [self]
+    def read_headers
+      @headers = @src_csv.shift
+      self
+    end
 
-  def additional_data(&block)
-    steps << [:additional_data, @headers, block]
-    self
-  end
+    # Computes a value once per batch, passed to the steps that follow as their third argument.
+    # @yieldparam batch [Array<Array<String>>]
+    # @yieldparam headers [Array<String>, nil]
+    # @return [self]
+    def additional_data(&block)
+      steps << [:additional_data, @headers, block]
+      self
+    end
 
-  def select(&block)
-    steps << [:select, @headers, block]
-    self
-  end
+    # Keeps the rows the block returns true for.
+    # @yieldparam row [Array<String>]
+    # @yieldparam headers [Array<String>, nil]
+    # @yieldparam additional_data [Object]
+    # @return [self]
+    def select(&block)
+      steps << [:select, @headers, block]
+      self
+    end
 
-  def reject(&block)
-    steps << [:reject, @headers, block]
-    self
-  end
+    # Drops the rows the block returns true for.
+    # @yieldparam (see #select)
+    # @return [self]
+    def reject(&block)
+      steps << [:reject, @headers, block]
+      self
+    end
 
-  def map(new_headers, &block)
-    steps << [:map, @headers, block]
-    @headers = new_headers
-    self
-  end
+    # Replaces each row with the block's result.
+    # @param new_headers [Array<String>, nil] the headers after this step
+    # @yieldparam (see #select)
+    # @yieldreturn [Array] the new row
+    # @return [self]
+    def map(new_headers, &block)
+      steps << [:map, @headers, block]
+      @headers = new_headers
+      self
+    end
 
-  def append(additional_headers, &block)
-    steps << [:append, @headers, block]
+    # Adds the block's values to the end of each row.
+    # @param additional_headers [Array<String>, nil] headers for the new values; nil drops the header row
+    # @yieldparam (see #select)
+    # @yieldreturn [Array] values to append
+    # @return [self]
+    def append(additional_headers, &block)
+      steps << [:append, @headers, block]
+      # without headers (read_headers not called) the output has no header row to extend
+      @headers = (@headers + additional_headers if @headers && additional_headers)
+      self
+    end
 
-    if additional_headers
-      @headers += additional_headers
-    else
-      @headers = nil
+    # Calls the block for each row without changing it.
+    # @yieldparam (see #select)
+    # @return [self]
+    def each(&block)
+      steps << [:each, @headers, block]
+      self
     end
 
-    self
-  end
+    # Replaces the headers written to the output.
+    # @param headers [Array<String>, nil]
+    # @return [self]
+    def set_headers(headers)
+      @headers = headers
+      self
+    end
 
-  def each(&block)
-    steps << [:each, @headers, block]
-    self
-  end
+    # Runs the steps over every row, writes the output and closes the files it opened, even when a step raises.
+    # @param batch_size [Integer] rows read and run through the steps at a time
+    # @return [void]
+    def process(batch_size = 10_000)
+      batch = []
 
-  def set_headers(headers)
-    @headers = headers
-    self
-  end
+      @dest_csv << @headers if @headers
 
-  def process(batch_size = 10_000, &block)
-    batch = []
+      steps_proc = proc do
+        steps.each do |step_type, current_headers, proc|
+          batch = process_step(step_type, current_headers, batch, &proc)
+        end
 
-    @dest_csv << @headers if @headers
+        batch.each { |row| @dest_csv << row }
 
-    steps_proc = Proc.new do
-      steps.each do |step_type, current_headers, proc|
-        batch = process_step(step_type, current_headers, batch, &proc)
+        batch = []
       end
 
-      batch.each { |row| @dest_csv << row }
+      while (row = @src_csv.shift)
+        batch << row
+        steps_proc.call if batch.size >= batch_size
+      end
 
-      batch = []
+      steps_proc.call if batch.size.positive?
+    ensure
+      @src_csv.close
+      @dest_csv.close
     end
-                   
-    while (row = @src_csv.shift)
-      batch << row
-      steps_proc.call if batch.size >= batch_size
+
+    private
+
+    def steps
+      @steps ||= []
     end
 
-    steps_proc.call if batch.size > 0
+    # each step type has a process_<type>_step method that changes the batch in place
+    def process_step(step_type, current_headers, batch, &)
+      send(:"process_#{step_type}_step", current_headers, batch, &)
+      batch
+    end
 
-    @src_csv.close
-    @dest_csv.close
-  end
+    def process_select_step(current_headers, batch)
+      batch.select! { |row| yield row, current_headers, @additional_data }
+    end
 
-  private
+    def process_reject_step(current_headers, batch)
+      batch.reject! { |row| yield row, current_headers, @additional_data }
+    end
 
-  def steps
-    @steps ||= []
-  end
+    def process_map_step(current_headers, batch)
+      batch.map! { |row| yield row, current_headers, @additional_data }
+    end
 
+    def process_append_step(current_headers, batch)
+      batch.map! { |row| row + yield(row, current_headers, @additional_data) }
+    end
 
-  def process_step(step_type, current_headers, batch, &block)
-    case step_type
-    when :select
-      batch.select! do |row|
-        block.call row, current_headers, @additional_data
-      end
-    when :reject
-      batch.reject! do |row|
-        block.call row, current_headers, @additional_data
-      end
-    when :map
-      batch.map! do |row|
-        block.call row, current_headers, @additional_data
-      end
-    when :append
-      batch.map! do |row|
-        row + block.call(row, current_headers, @additional_data)
-      end
-    when :additional_data
-      @additional_data = block.call(batch, current_headers)
-    when :each
-      batch.each do |row|
-        block.call(row, current_headers, @additional_data)
-      end
+    def process_additional_data_step(current_headers, batch)
+      @additional_data = yield batch, current_headers
     end
 
-    batch
+    def process_each_step(current_headers, batch)
+      batch.each { |row| yield row, current_headers, @additional_data }
+    end
   end
 end

data/lib/csv_utils/csv_wrapper.rb

@@ -1,51 +1,75 @@
-# Wraps a CSV object, if wrapper opens the csv file it will close it
-class CSVUtils::CSVWrapper
-  attr_reader :csv
+# frozen_string_literal: true
 
-  def initialize(csv, mode, csv_options)
-    open(csv, mode, csv_options)
-  end
+module CSVUtils
+  # Wraps a path or an open CSV so callers can read, write and close it the same way.
+  # It closes only the files it opened itself.
+  class CSVWrapper
+    # @return [CSV, #shift, #<<] the CSV being wrapped
+    attr_reader :csv
+
+    # @param csv [String, CSV, #shift, #<<] path to open, or a CSV (or array of rows) to use as is
+    # @param mode [String] file mode for a path
+    # @param csv_options [Hash] options passed to CSV.open for a path
+    def initialize(csv, mode, csv_options)
+      open(csv, mode, csv_options)
+    end
 
-  def self.open(file, mode, csv_options = {})
-    csv = new(file, mode, csv_options)
+    # Wraps a path or CSV; with a block, yields the wrapper and closes it afterwards, even when the block raises.
+    # @param file [String, CSV] path or CSV
+    # @param mode [String] file mode for a path
+    # @param csv_options [Hash] options passed to CSV.open for a path
+    # @yieldparam csv [CSVWrapper]
+    # @return [CSVWrapper, Object] the wrapper, or the block's value
+    def self.open(file, mode, csv_options = {})
+      csv = new(file, mode, csv_options)
+      return csv unless block_given?
 
-    if block_given?
-      yield csv
-      csv.close
-    else
-      csv
+      begin
+        yield csv
+      ensure
+        csv.close
+      end
     end
-  end
 
-  def open(csv, mode, csv_options)
-    if csv.is_a?(String)
-      @close_when_done = true
-      @csv = CSV.open(csv, mode, **csv_options)
-    else
-      @close_when_done = false
-      @csv = csv
+    # @api private
+    # @return [CSV, #shift, #<<]
+    def open(csv, mode, csv_options)
+      if csv.is_a?(String)
+        @close_when_done = true
+        @csv = CSV.open(csv, mode, **csv_options)
+      else
+        @close_when_done = false
+        @csv = csv
+      end
     end
-  end
 
-  def <<(row)
-    csv << row
-  end
+    # @param row [Array]
+    # @return [CSV, #<<]
+    def <<(row)
+      csv << row
+    end
 
-  def shift
-    csv.shift
-  end
+    # @return [Array<String>, nil] the next row, nil at the end
+    def shift
+      csv.shift
+    end
 
-  def rewind
-    csv.rewind
-  end
+    # Moves back to the first row.
+    # @return [void]
+    def rewind
+      csv.rewind
+    end
 
-  def close
-    csv.close if close_when_done?
-  end
+    # Closes the CSV if the wrapper opened it.
+    # @return [void]
+    def close
+      csv.close if close_when_done?
+    end
 
-  private
+    private
 
-  def close_when_done?
-    @close_when_done
+    def close_when_done?
+      @close_when_done
+    end
   end
 end

data/lib/csv_utils/version.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+module CSVUtils
+  # Gem version, bumped by release-please
+  VERSION = '0.6.0'
+end