Relearning Python #0: Introduction

#!/usr/bin/env ruby 

require 'find'
require 'fileutils'
require 'optparse'
require 'yaml'
require 'json'


# Default list of files to prune in search
PRUNE = ['.svn', 'CVS', 'CVSROOT', '.DS_Store', '.git']

# Unfortunate artifact of transition between Ruby 1.9 and 2.0
if self.class.const_defined?(:Encoding) then
    ENCODING_UTF8 = Encoding.find('UTF-8')
else
    ENCODING_UTF8 = nil
end

# Simple class to manage an ASCII spinner
class Spinner
    SPINNER_STATES = ['-', '\\', '|', "/"]

    SPINNER_INTERVAL = 0.5 # s

    def initialize(io=$stderr)
        @io = io
        @state = 0
        @updated = nil
    end

    def start()
        @updated = Time.now
        @io.print(SPINNER_STATES[0])
    end

    def update()
        if not @updated then
            start()
        end

        now = Time.now
        if now - @updated > SPINNER_INTERVAL then
            @updated = now
            @state = (@state + 1) % SPINNER_STATES.length
            @io.print("\b", SPINNER_STATES[@state])
        end
    end
end

# Collects duplicate file sets in lieu of an Array, 
# and writes them as Ruby objects
class DefaultPrinter
    attr_accessor :io, :pretty

    def initialize
        @results = []
    end

    def pretty?
        return @pretty
    end

    def << (obj)
        @results << obj
    end

    def print_results
        @results.sort!
        print(@results)
    end

    protected

    def print(obj)
        @io.write(obj)
        @io.write("\n");
    end
end

# Printer class (see above) that serializes as JSON
class JsonPrinter < DefaultPrinter
    def print(obj)
        if @pretty then
            text = JSON.pretty_generate(obj)
        else
            text = JSON.generate(obj)
        end
        @io.write(text)
        @io.write("\n")
    end
end

# Printer class (see above) that serializes as YAML.
class YamlPrinter < DefaultPrinter

    # :line_width setting so file names w/spaces don't break across lines
    DEFAULT_LINE_WIDTH = 4096

    def print(obj)
        opts = {:line_width => DEFAULT_LINE_WIDTH} 
        if @pretty then
            # {:canonical => true} looks almost identical to pretty-printed JSON
            # so the header keeps them distinct.
            opts[:header] = true
            opts[:canonical] = true
        end
        YAML.dump(@results, @io, opts)
    end
end


# Whether `path_i` and `path_j` refer to duplicate but not identical files
def duplicate_files?(path_i, path_j)
    return (File.exist?(path_i) \
            and File.exist?(path_j) \
            and not File.identical?(path_i, path_j) \
            and FileUtils.cmp(path_i, path_j))
end

def prunable?(path, prune=[])
    name = File.basename(path)
    return (prune.include?(name) or File.fnmatch('._*', name))
end

# Recurse through array of `dirs` and produce Hash of all files by size.
def files_by_size(dirs, prune=[], verbose=false)
    result = Hash.new {|h,k| h[k]=[]}
    count = 0
    spinner = nil 
    if verbose then
        $stderr.print("Looking for files in ", dirs, ": ")
        spinner = Spinner.new($stderr)
        spinner.start
    end
    Find.find(*dirs) do |path|
        if prunable?(path, prune) then
            Find.prune()
        elsif File.file?(path) then
            size = File.size(path)
            if size > 0 then
                count = count + 1
                path.encode!(ENCODING_UTF8) if ENCODING_UTF8
                result[size] << path if size > 0
                if verbose then
                    count = count + 1
                    spinner.update
                end
            end
        end
    end
    if verbose then
        $stderr.print(" done!\n")

        $stderr.print("Found ", count, " non-empty files in ", 
                      result.size, " size groups.\n")
    end
    return result
end

# Compare each file in `paths` and append lists of equal files to `result`
def append_duplicates(result, paths)
    idsets = Hash.new {|h,k| h[k]=[k]}
    0.upto(paths.length - 1) do |i|
        (i+1).upto(paths.length - 1) do |j|
            path_i = paths[i]
            path_j = paths[j]
            if duplicate_files?(path_i, path_j) then
                set_i, set_j = idsets[path_i], idsets[path_j]
                set_i << path_j
                set_i.sort!
                set_j << path_i
                set_j.sort!
            end
        end
    end
    idsets.values.uniq.each do | s |
        result << s
    end
    return result
end

# Find all files in `fmap` and append lists of equal files to `result`
# `fmap` is a map of files by size.
def append_all_dupes(result, fmap, verbose=false)
    fmap.each_pair do |size, paths|
        if paths.size > 1 then
            $stderr.print("Comparing ", paths, " ...") if verbose
            append_duplicates(result, paths)
            $stderr.print("done.\n") if verbose
        end
    end
end

# Find all files in `srcdir`, compare to all files in `fmap`,
# and append lists of equal files to `result`.
# The path from `srcdir` will always be first in the list.
def append_dir_dupes(result, fmap, srcdir, prune=[], verbose=false)
    $stderr.print("Looking for files in '", srcdir, "':\n") if verbose
    Find.find(srcdir) do |path_i|
        if prunable?(path_i, prune) then
            Find.prune()
        elsif File.file?(path_i) then
            size = File.size(path_i)
            paths = fmap[size]
            if size > 0 and paths and not paths.empty? then
                dupes = []
                $stderr.print("Comparing '", path_i, 
                              "' to ", paths, " ...") if verbose
                paths.each do |path_j|
                    if duplicate_files?(path_i, path_j) then
                        dupes << path_j
                    end
                end
                if dupes.size > 0 then
                   dupes.sort!
                   result << [path_i, *dupes]
                end
                $stderr.print(" done.\n") if verbose
            end
        end
    end
    $stderr.print("Done!\n") if verbose
    return result
end

def parse_options(config)
    
    outfile = nil
    pclass = YamlPrinter
    pretty = false

    OptionParser.new do |opts|
        opts.banner = "Usage: duplicate-files.rb [options] dir1 [dir2 ...]"

        opts.on("-q", "--[no-]quiet", "Run absolutely quietly") do |v|
            config.info = (not v)
        end
        opts.on("-v", "--[no-]verbose", "Run verbosely") do |v|
            config.verbose = v
        end
        opts.on("-d", "--from [DIR]", 
                "Compare arguments only to files from DIR") do |d|
            config.canondir = d
        end
        opts.on("-o", "--output [OUTFILE]", 
                "Write standard output to OUTFILE") do |f|
            outfile = File.new(f, 'w')
        end
        opts.on("-j", "--json", "Write output as JSON") do |v|
            pclass = JsonPrinter
        end
        opts.on("-y", "--yaml", "Write output as YAML (default)") do |v|
            pclass = YamlPrinter 
        end
        opts.on("-p", "--[no-]pretty", "Pretty-print output") do |v|
            pretty = v
        end
    end.parse!

    printer = pclass.new

    printer.io = (outfile or $stdout)
    printer.pretty = pretty

    config.printer = printer
    config.dirs = ARGV
end

Options = Struct.new("Options", 
                     :info, :verbose, :canondir, :dirs, :prune, :printer)

def run
    opts = Options.new(true, false, nil, [], PRUNE, nil)

    parse_options(opts)

    fmap = files_by_size(opts.dirs, opts.prune, (opts.info or opts.verbose))

    if opts.info then
        $stderr.print("Comparing files ...")
    end

    if opts.canondir then
        append_dir_dupes(opts.printer, fmap, opts.canondir, 
                         opts.prune, opts.verbose)
    else
        append_all_dupes(opts.printer, fmap, opts.verbose)
    end

    if opts.info then
        $stderr.print(" done!\n")
    end

    opts.printer.print_results
end

run()