sup-tweak-labels in sup-0.9

- old
+ new

@@ -1,9 +1,10 @@
 #!/usr/bin/env ruby
 
 require 'rubygems'
 require 'trollop'
+require 'enumerator'
 require "sup"
 
 class Float
   def to_s; sprintf '%.2f', self; end
    def to_time_s
@@ -35,12 +36,12 @@
 where <source>* is zero or more source URIs. Supported source URI schemes can
 be seen by running "sup-add --help".
 
 Options:
 EOS
-  opt :add, "One or more labels (comma-separated) to add to every message from the specified sources", :type => String
-  opt :remove, "One or more labels (comma-separated) to remove from every message from the specified sources, if those labels are present", :type => String
+  opt :add, "One or more labels (comma-separated) to add to every message from the specified sources", :default => ""
+  opt :remove, "One or more labels (comma-separated) to remove from every message from the specified sources, if those labels are present", :default => ""
   opt :query, "A Sup search query", :type => String
 
   text <<EOS
 
 Other options:
@@ -51,27 +52,28 @@
   opt :dry_run, "Don't actually modify the index. Probably only useful with --verbose.", :short => "-n"
   opt :version, "Show version information", :short => :none
 end
 opts[:verbose] = true if opts[:very_verbose]
 
-add_labels = (opts[:add] || "").split(",").map { |l| l.intern }.uniq
-remove_labels = (opts[:remove] || "").split(",").map { |l| l.intern }.uniq
+add_labels = opts[:add].to_set_of_symbols ","
+remove_labels = opts[:remove].to_set_of_symbols ","
 
 Trollop::die "nothing to do: no labels to add or remove" if add_labels.empty? && remove_labels.empty?
+Trollop::die "no sources specified" if ARGV.empty?
 
 Redwood::start
+index = Redwood::Index.init
+index.lock_interactively or exit
 begin
-  index = Redwood::Index.new
   index.load
 
-  source_ids = 
-    if opts[:all_sources]
-      index.sources
-    else
-      ARGV.map do |uri|
-        index.source_for uri or Trollop::die "Unknown source: #{uri}. Did you add it with sup-add first?"
-      end
+  source_ids = if opts[:all_sources]
+    Redwood::SourceManager.sources
+  else
+    ARGV.map do |uri|
+      Redwood::SourceManager.source_for uri or Trollop::die "Unknown source: #{uri}. Did you add it with sup-add first?"
+    end
   end.map { |s| s.id }
   Trollop::die "nothing to do: no sources" if source_ids.empty?
 
   query = "+(" + source_ids.map { |id| "source_id:#{id}" }.join(" OR ") + ")"
   if add_labels.empty?
@@ -79,33 +81,34 @@
     ## query to only messages with those labels
     query += " +(" + remove_labels.map { |l| "label:#{l}" }.join(" ") + ")"
   end
   query += ' ' + opts[:query] if opts[:query]
 
-  docs = Redwood::Index.run_query query
-  num_total = docs.size
+  parsed_query = index.parse_query query
+  parsed_query.merge! :load_spam => true, :load_deleted => true, :load_killed => true
+  ids = Enumerable::Enumerator.new(index, :each_id, parsed_query).map
+  num_total = ids.size
 
   $stderr.puts "Found #{num_total} documents across #{source_ids.length} sources. Scanning..."
 
   num_changed = num_scanned = 0
   last_info_time = start_time = Time.now
-  docs.each do |id|
+  ids.each do |id|
     num_scanned += 1
 
     m = index.build_message id
-    old_labels = m.labels.clone
+    old_labels = m.labels.dup
 
     m.labels += add_labels
     m.labels -= remove_labels
-    m.labels = m.labels.uniq
 
-    unless m.labels.sort_by { |s| s.to_s } == old_labels.sort_by { |s| s.to_s }
+    unless m.labels == old_labels
       num_changed += 1
       puts "From #{m.from}, subject: #{m.subj}" if opts[:very_verbose]
-      puts "#{m.id}: {#{old_labels.join ','}} => {#{m.labels.join ','}}" if opts[:verbose]
+      puts "#{m.id}: {#{old_labels.to_a.join ','}} => {#{m.labels.to_a.join ','}}" if opts[:verbose]
       puts if opts[:very_verbose]
-      index.sync_message m unless opts[:dry_run]
+      index.update_message_state m unless opts[:dry_run]
     end
 
     if Time.now - last_info_time > 60
       last_info_time = Time.now
       elapsed = last_info_time - start_time
@@ -116,10 +119,10 @@
   end
   $stderr.puts "Scanned #{num_scanned} / #{num_total} messages and changed #{num_changed}."
 
   unless num_changed == 0
     $stderr.puts "Optimizing index..."
-    index.ferret.optimize unless opts[:dry_run]
+    index.optimize unless opts[:dry_run]
   end
 
 rescue Exception => e
   File.open("sup-exception-log.txt", "w") { |f| f.puts e.backtrace }
   raise