feat: qtl xtrain for multiple train dirs

Author: c-dilks

bin/qtl-xtrain

@@ -7,67 +7,85 @@ unless ARGV.length == 2
   Verify that a directory of a train's skim files has the same list of
   run numbers as a directory of DST-file run directories.
 
-  USAGE: qtl xtrain [TRAIN_DIR] [DST_DIR]
+  USAGE: qtl xtrain [TRAIN_DIR(S)] [DST_DIR]
 
   Both directories must be on /mss
+
+  Use globs to specify multiple TRAIN_DIRS (surround in quotes)
   """
   exit 2
 end
-train_dir, dst_dir = ARGV
 
-# function to get a set of run numbers from one of the argument dirs
-def get_runnums(path, type)
-  runnums = Set.new
-  raise "#{type} dir `#{path}` is not on /mss" unless path.match? /^\/mss\//
-  raise "#{type} dir `#{path}` does not exist" unless Dir.exist? path
+train_dirs, dst_dir = ARGV
+num_failures = 0
 
-  # get list of files/directories within
-  files = []
-  case type
-  when :train
-    files = Dir.glob File.join(path, '*.hipo')
-  when :DST
-    files = Dir.glob File.join(path, '*/')
-  else
-    raise 'bad type'
-  end
-  raise "no #{type} files found in #{type} dir `#{path}`" if files.empty?
+# loop over train directories glob
+Dir.glob(train_dirs).each do |train_dir|
+
+  # function to get a set of run numbers from one of the argument dirs
+  def get_runnums(path, type)
+    runnums = Set.new
+    raise "#{type} dir `#{path}` is not on /mss" unless path.match? /^\/mss\//
+    raise "#{type} dir `#{path}` does not exist" unless Dir.exist? path
 
-  # extract their run numbers
-  files.each do |file|
-    nums = File.basename(file).scan(/\d+/).map &:to_i
-    raise "failed to get run number from #{type} object `#{file}`" unless nums.length == 1
-    runnums << nums[0]
+    # get list of files/directories within
+    files = []
+    case type
+    when :train
+      files = Dir.glob File.join(path, '*.hipo')
+    when :DST
+      files = Dir.glob File.join(path, '*/')
+    else
+      raise 'bad type'
+    end
+    raise "no #{type} files found in #{type} dir `#{path}`" if files.empty?
+
+    # extract their run numbers
+    files.each do |file|
+      nums = File.basename(file).scan(/\d+/).map &:to_i
+      raise "failed to get run number from #{type} object `#{file}`" if nums.empty?
+      runnums << nums[-1]
+    end
+    raise "failed to get run numbers from #{type} dir `#{path}`" if runnums.empty?
+    runnums
   end
-  raise "failed to get run numbers from #{type} dir `#{path}`" if runnums.empty?
-  runnums
-end
 
-# get runnum lists
-train_runs = get_runnums train_dir, :train
-dst_runs   = get_runnums dst_dir,   :DST
-puts """----------------------------------------------------------------------------------
-train dir run list:
-#{train_runs}
-DST dir run list:
-#{dst_runs}
-----------------------------------------------------------------------------------"""
+  # get runnum lists
+  train_runs = get_runnums train_dir, :train
+  dst_runs   = get_runnums dst_dir,   :DST
+  puts """
+  ==================================================================================
+  train dir: #{train_dir}
+  run list: #{train_runs}
+  -----
+  DST dir: #{dst_dir}
+  run list: #{dst_runs}
+  =================================================================================="""
 
-# compare runnum sets
-only_in_trains = train_runs - dst_runs
-only_in_dsts   = dst_runs   - train_runs
+  # compare runnum sets
+  only_in_trains = train_runs - dst_runs
+  only_in_dsts   = dst_runs   - train_runs
 
-# return results
-code = 0
-unless only_in_trains.empty?
-  $stderr.puts "ERROR: there are runs with skim files, but no corresponding DST-file directories:"
-  $stderr.puts only_in_trains
-  code = 1
-end
-unless only_in_dsts.empty?
-  $stderr.puts "ERROR: there are runs with DST-file directories, but no corresponding skim files:"
-  $stderr.puts only_in_dsts
-  code = 1
+  # print results
+  $stdout.flush
+  $stderr.flush
+  code = 0
+  unless only_in_trains.empty?
+    $stderr.puts "ERROR: there are runs with skim files, but no corresponding DST-file directories:"
+    $stderr.puts only_in_trains
+    code = 1
+  end
+  unless only_in_dsts.empty?
+    $stderr.puts "ERROR: there are runs with DST-file directories, but no corresponding skim files:"
+    $stderr.puts only_in_dsts
+    code = 1
+  end
+
+  # handle exit code
+  if code == 0
+    puts "All good"
+  else
+    num_failures += 1
+  end
 end
-puts "All good" if code == 0
-exit code
+exit num_failures > 0 ? 1 : 0