defrag.py - 'don't remember the original name/source

#!/usr/bin/env python
from __future__ import division
import sys, os, re
import getopt
import sys
import time
import os
import subprocess
from collections import deque

try:
    import psyco
    psyco.full()
except: pass
ext_pattern=re.compile(r'.*:\s+(\d+) extents? found')
def _supports_progress(f):
    if not hasattr(f, 'isatty'):
        return False
    if not f.isatty():
        return False
    if os.environ.get('TERM') == 'dumb':
        # e.g. emacs compile window
        return False
    return True

def ProgressBar(to_file=sys.stderr, **kwargs):
    """Abstract factory"""
    if _supports_progress(to_file):
        return TTYProgressBar(to_file=to_file, **kwargs)
    else:
        return DotsProgressBar(to_file=to_file, **kwargs)

class _BaseProgressBar(object):
    def __init__(self,
                 to_file=sys.stderr,
                 show_pct=False,
                 show_spinner=False,
                 show_bar=True,
                 show_count=True,
                 to_messages_file=sys.stdout,
                 _stack=None):
        object.__init__(self)
        self.to_file = to_file
        self.to_messages_file = to_messages_file
        self.last_msg = None
        self.last_cnt = None
        self.last_total = None
        self.show_pct = show_pct
        self.show_spinner = show_spinner
        self.show_bar = show_bar
        self.show_count = show_count
        self._stack = _stack
        # seed throttler
        self.MIN_PAUSE = 0.5 # seconds
        now = time.time()
        # starting now
        self.start_time = now
        # next update should not throttle
        self.last_update = now - self.MIN_PAUSE - 1

    def finished(self):
        """Return this bar to its progress stack."""
        self.clear()
        assert self._stack is not None
        self._stack.return_pb(self)

    def note(self, fmt_string, *args, **kwargs):
        """Record a note without disrupting the progress bar."""
        self.clear()
        self.to_messages_file.write(fmt_string % args)
        self.to_messages_file.write('\n')

    def child_progress(self, **kwargs):
        return ChildProgress(**kwargs)

class DotsProgressBar(_BaseProgressBar):

    def __init__(self, **kwargs):
        _BaseProgressBar.__init__(self, **kwargs)
        self.last_msg = None
        self.need_nl = False

    def tick(self):
        self.update()

    def update(self, msg=None, current_cnt=None, total_cnt=None):
        if msg and msg != self.last_msg:
            if self.need_nl:
                self.to_file.write('\n')

            self.to_file.write(msg + ': ')
            self.last_msg = msg
        self.need_nl = True
        self.to_file.write('.')

    def clear(self):
        if self.need_nl:
            self.to_file.write('\n')

    def child_update(self, message, current, total):
        self.tick()

class TTYProgressBar(_BaseProgressBar):
    """Progress bar display object.

    Several options are available to control the display.  These can
    be passed as parameters to the constructor or assigned at any time:

    show_pct
        Show percentage complete.
    show_spinner
        Show rotating baton.  This ticks over on every update even
        if the values don't change.
    show_eta
        Show predicted time-to-completion.
    show_bar
        Show bar graph.
    show_count
        Show numerical counts.

    The output file should be in line-buffered or unbuffered mode.
    """
    SPIN_CHARS = r'/-\|'


    def __init__(self, **kwargs):
        #from bzrlib.osutils import terminal_width
        #TODO: Determine terminal width
        _BaseProgressBar.__init__(self, **kwargs)
        self.spin_pos = 0
        self.width = 80
        self.start_time = time.time()
        self.last_updates = deque()
        self.child_fraction = 0


    def throttle(self):
        """Return True if the bar was updated too recently"""
        # time.time consistently takes 40/4000 ms = 0.01 ms.
        # but every single update to the pb invokes it.
        # so we use time.time which takes 20/4000 ms = 0.005ms
        # on the downside, time.time() appears to have approximately
        # 10ms granularity, so we treat a zero-time change as 'throttled.'

        now = time.time()
        interval = now - self.last_update
        # if interval > 0
        if interval < self.MIN_PAUSE:
            return True

        self.last_updates.append(now - self.last_update)
        self.last_update = now
        return False


    def tick(self):
        self.update(self.last_msg, self.last_cnt, self.last_total,
                    self.child_fraction)

    def child_update(self, message, current, total):
        if current is not None and total != 0:
            child_fraction = float(current) / total
            if self.last_cnt is None:
                pass
            elif self.last_cnt + child_fraction <= self.last_total:
                self.child_fraction = child_fraction
        if self.last_msg is None:
            self.last_msg = ''
        self.tick()


    def update(self, msg, current_cnt=None, total_cnt=None,
               child_fraction=0):
        """Update and redraw progress bar."""

        if current_cnt < 0:
            current_cnt = 0

        if current_cnt > total_cnt:
            total_cnt = current_cnt


        old_msg = self.last_msg
        # save these for the tick() function
        self.last_msg = msg
        self.last_cnt = current_cnt
        self.last_total = total_cnt
        self.child_fraction = child_fraction

        # each function call takes 20ms/4000 = 0.005 ms,
        # but multiple that by 4000 calls -> starts to cost.
        # so anything to make this function call faster
        # will improve base 'diff' time by up to 0.1 seconds.
        if old_msg == self.last_msg and self.throttle():
            return

        if self.show_spinner:
            spin_str = self.SPIN_CHARS[self.spin_pos % 4] + ' '
        else:
            spin_str = ''

        # always update this; it's also used for the bar
        self.spin_pos += 1

        if self.show_pct and self.last_total and self.last_cnt:
            pct = 100.0 * ((self.last_cnt + self.child_fraction) / self.last_total)
            pct_str = ' (%5.1f%%)' % pct
        else:
            pct_str = ''

        if not self.show_count:
            count_str = ''
        elif self.last_cnt is None:
            count_str = ''
        elif self.last_total is None:
            count_str = ' %i' % (self.last_cnt)
        else:
            # make both fields the same size
            t = '%i' % (self.last_total)
            c = '%*i' % (len(t), self.last_cnt)
            count_str = ' ' + c + '/' + t

        if self.show_bar:
            # progress bar, if present, soaks up all remaining space
            cols = self.width - 1 - len(self.last_msg) - len(spin_str) - len(pct_str) \
                   - len(count_str) - 3

            if self.last_total:
                # number of markers highlighted in bar
                markers = int(round(float(cols) *
                              (self.last_cnt + self.child_fraction) / self.last_total))
                bar_str = '[' + ('=' * markers).ljust(cols) + '] '
            elif False:
                # don't know total, so can't show completion.
                # so just show an expanded spinning thingy
                m = self.spin_pos % cols
                ms = (' ' * m + '*').ljust(cols)

                bar_str = '[' + ms + '] '
            else:
                bar_str = ''
        else:
            bar_str = ''

        m = spin_str + self.last_msg + bar_str + count_str + pct_str

        assert len(m) < self.width
        self.to_file.write('\r' + m.ljust(self.width - 1))
        #self.to_file.flush()

    def clear(self):
        self.to_file.write('\r%s\r' % (' ' * (self.width - 1)))
        #self.to_file.flush()


class ChildProgress(_BaseProgressBar):
    """A progress indicator that pushes its data to the parent"""
    def __init__(self, _stack, **kwargs):
        _BaseProgressBar.__init__(self, _stack=_stack, **kwargs)
        self.parent = _stack.top()
        self.current = None
        self.total = None
        self.child_fraction = 0
        self.message = None

    def update(self, msg, current_cnt=None, total_cnt=None):
        self.current = current_cnt
        self.total = total_cnt
        self.message = msg
        self.child_fraction = 0
        self.tick()

    def child_update(self, message, current, total):
        if current is None or total == 0:
            self.child_fraction = 0
        else:
            self.child_fraction = float(current) / total
        self.tick()

    def tick(self):
        if self.current is None:
            count = None
        else:
            count = self.current+self.child_fraction
            if count > self.total:
                count = self.total
        self.parent.child_update(self.message, count, self.total)

    def clear(self):
        pass

    def note(self, *args, **kwargs):
        self.parent.note(*args, **kwargs)


def str_tdelta(delt):
    if delt is None:
        return "-:--:--"
    delt = int(round(delt))
    return '%d:%02d:%02d' % (delt/3600,
                             (delt/60) % 60,
                             delt % 60)

def sort_by_value(d):
    """ Returns the keys of dictionary d sorted by their values """
    items=d.items()
    backitems=[ [v[1],v[0]] for v in items]
    backitems.sort()
    backitems.reverse()
    return [ backitems[i][1] for i in xrange(0,len(backitems))]

def copy(args):
    progress=subprocess.Popen(args,stdout=subprocess.PIPE,stderr=subprocess.PIPE).stdout
    c=progress.readline()
    c=progress.readline()
    c=progress.readline()
    while c != '':
        c=progress.read(1)
        if c =='\n':
            progress.read()
            sys.stdout.write(' '*80+'\r')
            return
        sys.stdout.write(c)
        if c == '\r':
            sys.stdout.write("     ")
            sys.stdout.flush()

def defrag(file, fs, fragments, threshold):
    #Performs a defrag
    if fragments<= threshold:
        print "     Already defragmented."
        return fragments
    status=subprocess.Popen(['lsof','-f','--',file], stdout=subprocess.PIPE).communicate()[0]
    if len(status)>1 and not status.startswith("COMMAND"):
        print "     Unable to determine if file is open, skipping..."
        return fragments

    if len(status)>1 and status.split('\n')[1].split()[3].find('w')>=0:
        print "     File is open for write, skipping..."
        return fragments
    import shutil
    try:
        os.mkdir(os.path.dirname(fs)+"/.defrag")
    except OSError: pass
    os.chmod(os.path.dirname(fs)+"/.defrag",0)
    old_mtime=os.path.getmtime(file)
    copy(["rsync",'-a','--progress', file,os.path.dirname(fs)+"/.defrag/"])
    new_numfrags=numfrags(os.path.dirname(fs)+"/.defrag/"+os.path.basename(file))
    if new_numfrags >= fragments:
        print "     No improvement (%.1f --> %.1f)" % (fragments,new_numfrags)
        os.unlink(os.path.dirname(fs)+"/.defrag/"+os.path.basename(file))
        return fragments
    else:
        if new_numfrags <= threshold:
            print "     Fully defragmented!"
        else:
            print "     Improved: (%.1f --> %.1f)" % (fragments,new_numfrags)

        if os.path.getmtime(file) == old_mtime:
            shutil.move(os.path.dirname(fs)+"/.defrag/"+os.path.basename(file), file)
            return new_numfrags
        else:
            print "     Aborted: file changed during defrag."
            return numfrags(file)

def build_filelist(fs):
    #Uses find to get list of files on desired filesystem
    list=subprocess.Popen(["/usr/bin/find",fs,'-xdev','-type','f', '-print0'],stdout=subprocess.PIPE).communicate()[0].split('\0')
    return list[:-1]
def numfrags(file):
    #Uses filefrag to determine # of fragments.
    fragresult=subprocess.Popen(["filefrag",file],stdout=subprocess.PIPE).communicate()[0]
    match=ext_pattern.search(fragresult)
    size = os.path.getsize(file)
    if not match:
        print "Error analyzing",file
        return 0
    else:
        frags=int(match.expand(r'\1'))
        if frags <= 1 or size == 0:
            return 0
    return (frags/(size/unit))


def run(path,threshold=0.5,passes=-1,list=None, fragmatrix=None):
    if not fragmatrix:
        print "Building list of files to analyze...",
        sys.stdout.flush()
        list=build_filelist(path)
        print "done!"
        fragmatrix={}
        n=0
        pb = ProgressBar(show_pct=True, show_bar=True, show_spinner=True)
        for file in list:
            pb.update((file+" "*25)[:25]+" ",float(n),len(list))
            n+=1
            f=numfrags(file)
            if f> threshold: fragmatrix[file]=f
        pb.clear()
        print "\nAnalyze finished.",
        if len(fragmatrix) == 0:
            print "\n%s is not fragmented. Go home." % path
            sys.exit(0)
    most_fragmented_files=sort_by_value(fragmatrix)
    frags=sum(fragmatrix.values())
    print "%.1f" % (len(most_fragmented_files)*100.0/len(list)),"%% fragmentation (%d files)," % len(most_fragmented_files),  "%.1f average frags" % (float(frags)/len(fragmatrix))
    print "Fragmented files:"
    try:
        for i in xrange(0,10):
            print "%.2f\t%s" % (fragmatrix[most_fragmented_files[i]], most_fragmented_files[i])
        print "..."
    except IndexError: pass
    if passes <0:
        while True:
            try:
                tmp=raw_input("How many passes to run? [10] ")
                if len(tmp)==0:
                    passes=-10
                    break
                passes=int(tmp)*abs(passes)/passes
                break
            except ValueError: pass
    try:
        for k in xrange(0,abs(passes)):
            n=0
            rfrags=0
            t=0
            for t in most_fragmented_files:
                rfrags+=fragmatrix[t]
            print "\n",(str(k+1))*60
            print "===> Pass", k+1,"of", abs(passes),"<===    Remaining Fragmentation %d/%d (%d%%)" % (rfrags,frags,(rfrags/frags)*100)
            print (str(k+1))*60+"\n"
            for i in most_fragmented_files:
                n+=1
                print "\n  Pass %d of %d, %d/%d (%d%%):" % (k+1,abs(passes),n,len(most_fragmented_files),100*n/len(most_fragmented_files)), "%.1f frags" % fragmatrix[i], i
                try:
                    fragmatrix[i]=defrag(i,path,fragmatrix[i],threshold)
                except OSError:
                    print "     Error defragmenting. Did the file disappear?"
            tmp=most_fragmented_files[:]
            for i in tmp:
            if fragmatrix[i] <= threshold:
                    most_fragmented_files.remove(i)

    finally:

        print "========= COMPLETE ==========="
        print "Remaining Fragmented Files:"
        try:
            most_fragmented_files.reverse()
            for i in most_fragmented_files:
                print "%.2f\t%s" % (fragmatrix[i], i)
        except IndexError: pass
        print "Frags/unit Before:\t %.2f" % frags
        nowfrags=0
        for i in most_fragmented_files:
            nowfrags+=fragmatrix[i]
        print "Frags/unit After: \t %.2f" % nowfrags
        print "Improvement:    \t %.1f %%" % abs((frags-nowfrags)/frags * 100)
        print "==============================="
        if nowfrags and passes < 0:
            response='Fish'
            while not response in ('Y','y','N','n',''):
                response=raw_input("\nThere is still fragmentation. Run another set of passes? [Y/n]")
            if response in ('Y','y',''):
                run(path,threshold,passes,list,fragmatrix)

def fillunits():
    udict={}
    n = 1
    ni = 1
    for key in ('B','K','M','G','T','P','E','Z','Y'):
    udict[key]=n
    udict[key+'B']=n
    udict[key+'I']=ni
    udict[key+'IB']=ni
    n *= 1000
    ni *= 1024
    return udict

try:
    opts=None
    try:
        opts=getopt.gnu_getopt(sys.argv[1:],'hn:t:a:u',('help', 'passes=', 'threshold=' ,'analyze','unit='))
        passes = -1
        threshold = .5
        unit = 1024*1024
        unitstr='MiB'
        unitdict=fillunits()
        for i in opts[0]:
            if   i[0] in ('-n','--passes'):
                passes=int(i[1])
            elif i[0] in ('-h', '--help'):
                raise getopt.GetoptError("")
            elif i[0] in ('-t','--threshold'):
        try:
            threshold=float(i[1])
        except:
            raise getopt.GetoptError("Threshold must be a float")
            elif i[0] in ('-a','--analyze'):
                passes=0
                threshold=0
            elif i[0] in ('-u','--unit'):
        if i[1].isdigit(): unit = int(i[1])
        elif i[1].upper() in unitdict.keys():
            unit = unitdict[i[1].upper()]
            unitstr=i[1]
        else: raise getopt.GetoptError("Unit may be a positive integer or a unit (B KB MB GB TB PB EB ZB YB):\n with K,KB=1000, KI=KIB=1024 ,etc")
    del unitdict
    if os.getuid():
        if passes == 0:
        print "This usually requires root access, but will try anyway"
        else:
        raise getopt.GetoptError("Everything bu analysis always requires root/sudo access.")
        try:
            if os.path.isdir(opts[1][0]) and opts[1][0][-1] != '/':
                opts[1][0]+="/"
            run(opts[1][0], threshold, passes)
        except IndexError:
            raise getopt.GetoptError("No path specified")
    except getopt.GetoptError, info:
        print info
        print "%s [-h] [-n passes] [--passes n] [-t threshold] [--threshold n] [-a] [--analyze] [-u unit] [--unit u] [--help] path"% sys.argv[0]
finally:
    if opts and len(opts[1]) and os.path.isdir(os.path.dirname(opts[1][0])+"/.defrag"):
        subprocess.Popen(['rm','-f','-r',os.path.dirname(opts[1][0])+"/.defrag"], stdout=subprocess.PIPE).communicate()[0]