122 lines
3.0 KiB
Ruby
122 lines
3.0 KiB
Ruby
require 'csv'
|
|
require 'json'
|
|
require_relative 'storage_plugin'
|
|
|
|
class DataPipeline
|
|
#
|
|
# options
|
|
# file: path of the csv file
|
|
# csv: options passed to ruby CSV, merged with DEFAULTS
|
|
# columns: fixed column order for writing (and keys for reading without headers)
|
|
# without columns the values are written in hash order
|
|
# layout: DataPipeline::Layout, field order per record type (several record types in one file),
|
|
# unknown fields raise; no header line
|
|
# counter: print a dot per written row
|
|
# verbose: print settings
|
|
#
|
|
class CsvStorage < StoragePlugin
|
|
DEFAULTS = {
|
|
verbose: false,
|
|
counter: false,
|
|
columns: nil,
|
|
csv: {
|
|
headers: true,
|
|
col_sep: ',',
|
|
empty_value: ''
|
|
}
|
|
}.freeze
|
|
|
|
def initialize **options, &block
|
|
@options = DEFAULTS.merge(options) do |_key, default, value|
|
|
default.is_a?(Hash) && value.is_a?(Hash) ? default.merge(value) : value
|
|
end
|
|
@file_name = @options.fetch(:file) { raise ArgumentError, 'CsvStorage: option file missing' }
|
|
@verbose = @options[:verbose]
|
|
|
|
if @verbose
|
|
puts "CsvStorage#options--------"
|
|
puts JSON.pretty_generate @options
|
|
puts "---------------------------"
|
|
end
|
|
|
|
yield if block_given?
|
|
end
|
|
|
|
def read(&block)
|
|
if block_given?
|
|
CSV.foreach(@file_name, **csv_options) { |row| yield to_hash(row) }
|
|
return
|
|
end
|
|
|
|
CSV.read(@file_name, **csv_options).map { |row| to_hash(row) }
|
|
end
|
|
|
|
def write data
|
|
puts "> CsvStorage#write" if @verbose
|
|
open_file
|
|
data.each { |row| write_row row }
|
|
ensure
|
|
close
|
|
end
|
|
|
|
def write_row row
|
|
open_file
|
|
if @options[:layout]
|
|
@file << @options[:layout].values(row)
|
|
print '.' if @options[:counter]
|
|
return
|
|
end
|
|
|
|
if !@header_written && csv_options[:headers]
|
|
header = columns || row.keys
|
|
@file << header.map(&:to_s) if header
|
|
end
|
|
@header_written = true
|
|
|
|
@file << (columns ? columns.map { |col| fetch(row, col) } : row.values)
|
|
print "." if @options[:counter]
|
|
end
|
|
|
|
def close
|
|
@file&.close
|
|
@file = nil
|
|
@header_written = false
|
|
end
|
|
|
|
private
|
|
|
|
def csv_options
|
|
@options[:csv]
|
|
end
|
|
|
|
def columns
|
|
@options[:columns]
|
|
end
|
|
|
|
def open_file
|
|
return if @file
|
|
|
|
@file = CSV.open(@file_name, 'w', **csv_options)
|
|
@header_written = false
|
|
return unless columns && csv_options[:headers] && !@options[:layout]
|
|
|
|
@file << columns.map(&:to_s)
|
|
@header_written = true
|
|
end
|
|
|
|
# rows with headers are CSV::Row, without headers plain arrays
|
|
def to_hash row
|
|
return row.to_h if row.is_a?(CSV::Row)
|
|
|
|
keys = columns&.map(&:to_s) || (0...row.length).map(&:to_s)
|
|
keys.zip(row).to_h
|
|
end
|
|
|
|
def fetch row, col
|
|
return row[col] if row.key?(col)
|
|
|
|
row.key?(col.to_s) ? row[col.to_s] : row[col.to_s.to_sym]
|
|
end
|
|
end
|
|
end
|