1#!/usr/bin/env ruby
2# Sizes - Calculate and sort all filesizes for current folder
3# Includes directory sizes, colorized output
4# Brett Terpstra 2019 WTF License
5VERSION = "1.0.0"
6
7require 'shellwords'
8
9# Just including term-ansicolor by @flori and avoiding all the
10# rigamarole of requiring multiple files when it's not a gem... - Brett
11#
12# ansicolor Copyright: Florian Frank
13# License: <https://github.com/flori/term-ansicolor/blob/master/COPYING>
14# Home: <https://github.com/flori/term-ansicolor>
15module Term
16
17  # The ANSIColor module can be used for namespacing and mixed into your own
18  # classes.
19  module ANSIColor
20    # require 'term/ansicolor/version'
21
22    # :stopdoc:
23    ATTRIBUTES = [
24      [ :clear              ,   0 ],     # String#clear is already used to empty string in Ruby 1.9
25      [ :reset              ,   0 ],     # synonym for :clear
26      [ :bold               ,   1 ],
27      [ :dark               ,   2 ],
28      [ :italic             ,   3 ],     # not widely implemented
29      [ :underline          ,   4 ],
30      [ :underscore         ,   4 ],     # synonym for :underline
31      [ :blink              ,   5 ],
32      [ :rapid_blink        ,   6 ],     # not widely implemented
33      [ :negative           ,   7 ],     # no reverse because of String#reverse
34      [ :concealed          ,   8 ],
35      [ :strikethrough      ,   9 ],     # not widely implemented
36      [ :black              ,  30 ],
37      [ :red                ,  31 ],
38      [ :green              ,  32 ],
39      [ :yellow             ,  33 ],
40      [ :blue               ,  34 ],
41      [ :magenta            ,  35 ],
42      [ :cyan               ,  36 ],
43      [ :white              ,  37 ],
44      [ :on_black           ,  40 ],
45      [ :on_red             ,  41 ],
46      [ :on_green           ,  42 ],
47      [ :on_yellow          ,  43 ],
48      [ :on_blue            ,  44 ],
49      [ :on_magenta         ,  45 ],
50      [ :on_cyan            ,  46 ],
51      [ :on_white           ,  47 ],
52      [ :intense_black      ,  90 ],    # High intensity, aixterm (works in OS X)
53      [ :intense_red        ,  91 ],
54      [ :intense_green      ,  92 ],
55      [ :intense_yellow     ,  93 ],
56      [ :intense_blue       ,  94 ],
57      [ :intense_magenta    ,  95 ],
58      [ :intense_cyan       ,  96 ],
59      [ :intense_white      ,  97 ],
60      [ :on_intense_black   , 100 ],    # High intensity background, aixterm (works in OS X)
61      [ :on_intense_red     , 101 ],
62      [ :on_intense_green   , 102 ],
63      [ :on_intense_yellow  , 103 ],
64      [ :on_intense_blue    , 104 ],
65      [ :on_intense_magenta , 105 ],
66      [ :on_intense_cyan    , 106 ],
67      [ :on_intense_white   , 107 ]
68    ]
69
70    ATTRIBUTE_NAMES = ATTRIBUTES.transpose.first
71    # :startdoc:
72
73    # Returns true if Term::ANSIColor supports the +feature+.
74    #
75    # The feature :clear, that is mixing the clear color attribute into String,
76    # is only supported on ruby implementations, that do *not* already
77    # implement the String#clear method. It's better to use the reset color
78    # attribute instead.
79    def support?(feature)
80      case feature
81      when :clear
82        !String.instance_methods(false).map(&:to_sym).include?(:clear)
83      end
84    end
85    # Returns true, if the coloring function of this module
86    # is switched on, false otherwise.
87    def self.coloring?
88      @coloring
89    end
90
91    # Turns the coloring on or off globally, so you can easily do
92    # this for example:
93    #  Term::ANSIColor::coloring = STDOUT.isatty
94    def self.coloring=(val)
95      @coloring = val
96    end
97    self.coloring = true
98
99    ATTRIBUTES.each do |c, v|
100      eval <<-EOT
101        def #{c}(string = nil)
102          result = ''
103          result << "\e[#{v}m" if Term::ANSIColor.coloring?
104          if block_given?
105            result << yield
106          elsif string.respond_to?(:to_str)
107            result << string.to_str
108          elsif respond_to?(:to_str)
109            result << to_str
110          else
111            return result #only switch on
112          end
113          result << "\e[0m" if Term::ANSIColor.coloring?
114          result
115        end
116      EOT
117    end
118
119    # Regular expression that is used to scan for ANSI-sequences while
120    # uncoloring strings.
121    COLORED_REGEXP = /\e\[(?:(?:[349]|10)[0-7]|[0-9])?m/
122
123    # Returns an uncolored version of the string, that is all
124    # ANSI-sequences are stripped from the string.
125    def uncolored(string = nil) # :yields:
126      if block_given?
127        yield.to_str.gsub(COLORED_REGEXP, '')
128      elsif string.respond_to?(:to_str)
129        string.to_str.gsub(COLORED_REGEXP, '')
130      elsif respond_to?(:to_str)
131        to_str.gsub(COLORED_REGEXP, '')
132      else
133        ''
134      end
135    end
136
137    module_function
138
139    # Returns an array of all Term::ANSIColor attributes as symbols.
140    def attributes
141      ATTRIBUTE_NAMES
142    end
143    extend self
144  end
145end
146
147# Begin sizes
148
149class String
150
151  include Term::ANSIColor
152
153  # ensure trailing slash
154  def slashit
155    self.sub(/\/?$/,'/')
156  end
157
158  # colorize a human readable size format by size
159  def color_fmt
160    case self
161    when /\dB?$/
162      self.blue
163    when /\dKB?$/
164      self.green
165    when /\dMB?$/
166      self.yellow
167    when /\dGB?$/
168      self.red
169    else
170      self.bold.red
171    end
172  end
173
174  # colorize files by type (directories and hidden files)
175  def color_file(force_check=false)
176    filename = self.dup
177    if force_check && File.directory?(filename)
178      filename.sub!(/\/?$/,'/')
179    end
180
181    case filename
182    when /\/$/
183      filename.green
184    when /^\./
185      filename.white
186    else
187      filename.bold.white
188    end
189  end
190
191  # Replace $HOME in path with ~
192  def short_dir
193    home = ENV['HOME']
194    self.sub(/#{home}/, '~')
195  end
196
197  # Convert a line like `120414 filename` to a colorized string with
198  # human readable size
199  def line_to_human
200    parts = self.split(/\t/)
201    if parts[0] =~ /NO ACCESS/
202      "  ERROR".red + " " + parts[1].color_file
203    else
204      size = to_human(parts[0].to_i).color_fmt
205      size.pad_escaped(7) + " " + parts[1].color_file
206    end
207  end
208
209  # Pad a line containing ansi escape codes to a given length, ignoring
210  # the escape codes
211  def pad_escaped(len)
212    str = self.dup
213    str.gsub!(/\e\[\d+m/,'')
214    prefix = ""
215    while prefix.length + str.length < len
216      prefix += " "
217    end
218    prefix + self
219  end
220end
221
222# Convert a number (assumed bytes) to a human readable format (12.5K)
223def to_human(n,fmt=false)
224  count = 0
225  formats = %w(B K M G T P E Z Y)
226
227  while  (fmt || n >= 1024) && count < 8
228    n /= 1024.0
229    count += 1
230    break if fmt && formats[count][0].upcase =~ /#{fmt[0].upcase}/
231  end
232
233  format("%.2f%s",n,formats[count])
234end
235
236# Use `du` to size a single directory and all of its contents. This
237# number is returned in blocks (512B), so the human readable result may
238# be slightly different than you'd get from `ls` or a GUI file manager
239def du_size_single(dir)
240  res = %x{du -s #{Shellwords.escape(dir)} 2>/dev/null}.strip
241  if $?.success?
242    parts = res.split(/\t/)
243    {:size => parts[0].to_i * 512, :file => parts[1].strip}
244  else
245    {:size => nil, :file => dir}
246  end
247end
248
249
250# main function
251def all_sizes(dir)
252  # Use `ls` to list all files in the target with long info
253  files = %x{ls -lSrAF #{dir.slashit} 2>/dev/null}
254  unless $?.success?
255    $stdout.puts "Error getting file listing".red
256    Process.exit 1
257  end
258  files = files.strip.split(/\n/)
259
260  files.delete_if {|line|
261    line.strip =~ /^total \d+/
262  }
263
264  # trim file list to just size and filename
265  matches = files.map {|line|
266    # Examples of output I'm trying to match:
267    #
268    #     -rw-r--r--   1 alexwlchan  staff   1053  1 Jul  2018 LICENSE
269    #     -rw-r--r--   1 alexwlchan  staff    306  8 Jan 20:36 Makefile
270    #
271    m = line.match(/^\S{10,11} +\d+ +\S+ +\w+ +(?<size>\d+) +\d+ +\w{3} +[\d:]+ +(?<file>.*?)$/)
272    {:size => m[:size].to_i, :file => m[:file]}
273  }
274
275
276  # if a line is a path to a directory, use `du` to update its size with
277  # the total filesize of the directory contents.
278  matches.map! {|match|
279    if File.directory?(match[:file])
280      du_size_single(match[:file])
281    else
282      match
283    end
284  }
285
286  # Sort by size (after updating directory sizes)
287  matches.sort! {|a,b|
288    size1 = a[:size]
289    size2 = b[:size]
290    size1 <=> size2
291  }
292
293  # Output each line with human-readable size and colorization
294  matches.each {|match|
295    entry = "#{match[:size]}\t#{match[:file]}"
296    $stdout.puts entry.line_to_human
297  }
298  # Include a total for the directory
299  $stdout.puts "-------".black.bold
300  total = du_size_single(dir)
301  entry = "#{total[:size]}\t#{total[:file]}"
302  $stdout.puts(entry.short_dir.line_to_human)
303end
304
305def help
306  app = File.basename(__FILE__)
307  help =<<EOHELP
308#{app.bold.white} #{VERSION.green} by Brett Terpstra
309  Display a human-readable list of sizes for all files and directories.
310usage:
311  $ #{app.bold.white} [directory]
312Leaving directory blank operates in the current working directory.
313EOHELP
314  puts help
315  Process.exit 0
316end
317
318# Assume operating on current directory...
319dir = ENV['PWD']
320
321# ...unless an argument is provided
322if ARGV[0]
323  # Add some help. Why not?
324  if ARGV[0] =~ /^-?-h(elp)?/
325    help
326  elsif ARGV[0] =~ /^-?-v(ersion)?/
327    $stdout.puts File.basename(__FILE__) + " v" + VERSION
328    Process.exit 0
329  else
330    argdir = File.expand_path(ARGV[0])
331    if File.directory?(argdir)
332      dir = argdir
333    end
334  end
335end
336
337all_sizes(dir)