File usage stats in ruby

dirstat.rb

#!/usr/bin/ruby
# Dirstat by Sunny Ripert <sripert@insia.org>
# Provides useful file-usage statistics on a given directory.
#
# Example:
#
#   $ ruby dirstat.rb .
#           directory          %use       files           %use(graph)       size  
#   --------------------------------------------------------------------------------
#   ./Futurama               56.11 % (     27 files) |###########        |   4.58GB
#   ./IT-Crowd               22.26 % (      8 files) |####               |   1.82GB
#   --------------------------------------------------------------------------------

class String
  # ANSI-colored version of the string
  def colorize(color)
    n = case color
      when :red: 9
      when :orange: 1
      when :yellow: 11
      else return self
    end
    "\x1b[38;5;#{n}m#{self}\x1b[0m"
  end
end


# A DirStat object represents a file-system directory and provides stats methods.
class DirStat
  attr_reader :path

  def initialize(path)
    @path = path.gsub /(.)\/$/, '\1' # remove trailing slash except for '/'
  end

  # Array of DirStat objects of all the direct children directories, 
  # sorted by size
  def child_directories
    Dir.entries(@path).select { |path|
      path != '.' and path != '..' and File.directory?(File.join(@path, path))
    }.collect { |dir|
      DirStat.new(File.join(@path, dir))
    }.sort_by { |dir| -dir.size }
  end

  # Array of all descending files (cached)
  def descendant_files
    return @files unless @files.nil?
    pattern = File.join(@path, '**', '**')
    @files = Dir.glob(pattern, File::FNM_DOTMATCH).select { |p| File.file? p }
  end

  # Total size of the directory (cached)
  def size
    return @size unless @size.nil?
    @size = descendant_files.inject(0) { |acc, f| acc + File.size(f) }
  end

  # Total size of the directory in human format (B, kB, MB or GB)
  def human_size
    case
      when size < 1024:    '%.2f B' % size
      when size < 1024**2: '%.2fkB' % (size / 1024.0)
      when size < 1024**3: '%.2fMB' % (size / 1024.0**2)
      else                 '%.2fGB' % (size / 1024.0**3)
    end
  end

  # Print full stats on the directory. The print_color argument turns on
  # ANSI colors
  def show(print_color = false)
    puts 'directory'.center(25) + '%use'.center(8) + 'files'.center(16) + \
      '%use(graph)'.center(21) + 'size'.center(8)
    puts '-' * 80
    child_directories.each do |f|
      ratio = f.size/size.to_f
      color = case
        when !print_color: nil
        when ratio > 0.75: :red
        when ratio > 0.5:  :orange
        when ratio > 0.25: :yellow
        else               nil
      end
      printf "%-25s%5.2f %% (%7d file%s) |%s| %8s\n",
        f.path,
        ratio * 100,
        f.descendant_files.size,
        f.descendant_files.size == 1 ? ' ' : 's',
        ('#' * (ratio * 20)).ljust(20).colorize(color),
        f.human_size
    end
    puts '-' * 80
  end
end

if $0 == __FILE__
  dir = '.'
  color = true

  ARGV.each do |arg|
    if arg == '--nocolor'
      color = false
    elsif not File.directory?(arg)
      puts 'usage: #{$0} [options] [path]'
      puts 'options:'
      puts '	--nocolor      : disable output coloring'
      abort
    else
      dir = arg
    end
  end

  DirStat.new(dir).show(color)
end