pysize.py

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