pysize.py

   1 #!/usr/bin/env python
   2
   3 import os
   4 import optparse
   5
   6 from ui.ascii import ui_ascii
   7 from ui.curses import ui_curses
   8 from ui.gtk import ui_gtk
   9 from core.sigquit_traceback import install_sigquit_traceback
  10
  11 def _ui_auto(options, args):
  12     """Automatically choose the best available UI."""
  13     if ui_gtk.is_available():
  14         return ui_gtk.run(options, args)
  15
  16     if ui_curses.is_available():
  17         return ui_curses.run(options, args)
  18
  19     if ui_ascii.is_available():
  20         return ui_ascii.run(options, args)
  21
  22     print 'No UI available'
  23
  24 UI = {'ascii': ui_ascii.run, 'curses': ui_curses.run,
  25       'gtk': ui_gtk.run, 'auto': _ui_auto}
  26
  27 def _profile(continuation):
  28     try:
  29         import cProfile
  30         import pstats
  31         print 'Profiling using cProfile'
  32         cProfile.runctx('continuation()', globals(), locals(), 'pysize.prof')
  33         stats = pstats.Stats('pysize.prof')
  34     except ImportError:
  35         import hotshot
  36         import hotshot.stats
  37         prof = hotshot.Profile('pysize.prof')
  38         print 'Profiling using hotshot'
  39         prof.runcall(continuation)
  40         prof.close()
  41         stats = hotshot.stats.load('pysize.prof')
  42     stats.strip_dirs()
  43     stats.sort_stats('time', 'calls')
  44     stats.print_stats(40)
  45
  46 def _main():
  47     install_sigquit_traceback()
  48     parser = optparse.OptionParser(usage='pysize [OPTIONS] [DIRECTORY]')
  49     parser.add_option('--ui', type='choice', choices=UI.keys(), default='auto')
  50     parser.add_option('--max-depth', type='int', dest='max_depth', default=5)
  51     parser.add_option('--min-size', type='str', dest='min_size', default='auto')
  52     parser.add_option('--profile', action='store_true', dest='profile',
  53                       default=False)
  54     (options, args) = parser.parse_args()
  55
  56     if len(args) == 0:
  57         args = ['.']
  58     elif len(args) == 1:
  59         args[0] = os.path.realpath(args[0])
  60     else:
  61         print parser.error('At most a single argument is expected')
  62
  63     continuation = lambda: UI[options.ui](options, args)
  64
  65     if options.profile:
  66         _profile(continuation)
  67     else:
  68         continuation()
  69
  70 if __name__ == '__main__':
  71     try:
  72         # Try to use psyco if available
  73         import psyco
  74         psyco.full()
  75     except ImportError:
  76         pass
  77     _main()