sep.py

#!/bin/env python
# https://gitorious.org/mles/mles/source/5c575ae60df0b483d079120c17881d29c3bca7b3:sep.py
# http://grissiom.blogspot.in

import re
import os
from os import listdir, makedirs, rmdir, remove
from os.path import isdir, isfile, islink
joinp, splitp = os.path.join, os.path.split
from shutil import move, copy2
import sys

def exit_err(no, msg, stream=sys.stderr):
        sys.stdout.flush()
        stream.write('\n' + get_color('red') +
                         msg + get_color('default') + '\n')
        sys.exit(no)

def printf(arg, lv=0, stream=sys.stdout):
        if lv <= PRINT_LV:
                stream.write(arg)

def get_color(cl):
        return colours.get(cl, '')

# adapted from /usr/lib64/python2.6/os.py
def walk_dir(top, topdown=True, followlinks=False, ml=0, rl=0):
        '''walk like os.walk

        ml is the max recursion level. Recursion check is disabled if ml == 0.
        rl is the current recursion level. You don't need to set this variable
            at normal cases.'''
        if ml != 0 and rl >= ml:
                return
        try:
                names = listdir(top)
        except:
                pass

        dirs, nondirs = [], []
        for name in names:
                if isdir(joinp(top, name)):
                        dirs.append(name)
                else:
                        nondirs.append(name)

        if topdown:
                yield top, dirs, nondirs
        for name in dirs:
                path = joinp(top, name)
                if followlinks or not islink(path):
                        for x in walk_dir(path, topdown, followlinks, rl + 1):
                                yield x
        if not topdown:
                yield top, dirs, nondirs

def movefiles(file_name, destdir):
        if isdir(destdir) is False:
                try:
                        makedirs(destdir)
                except OSError, oserr:
                        exit_err(3, 'Error when creating %s: %s' %
                                        (oserr.filename, oserr.strerror))
        try:
                printf(get_color('yellow') + OPER_FILE_MOD
                      + ' ' + file_name + ' to '
                      + destdir + os.sep + ' ... ')
                sys.stdout.flush()
                # work around: shutil.move will fail if the dest file exits...
                if OPERATE_FILE == move:
                        if isfile(joinp(destdir, splitp(file_name)[1])):
                                remove(joinp(destdir, splitp(file_name)[1]))
                OPERATE_FILE(file_name, destdir)
                printf(get_color('green') + "Done\n" + get_color('default'))
        except IOError, ioerr:
                if isfile(joinp(destdir, splitp(file_name)[1])):
                        remove(joinp(destdir, splitp(file_name)[1]))
                exit_err(3, 'Error when %s %s: %s' %
                                (OPER_FILE_MOD, file_name, ioerr.strerror))
        except KeyboardInterrupt:
                if isfile(joinp(destdir, splitp(file_name)[1])):
                        remove(joinp(destdir, splitp(file_name)[1]))
                printf("Remove " + joinp(destdir, file_name) + '\n')
                exit_err(1, 'User aborted.')

def main(sourcedir, destdir):
        try:
                for root_dir, dirs, files in walk_dir(sourcedir, topdown=False, ml=RECURSION_LEVEL):
                        files = [ i for i in files if FILE_FILTER.search(i)]
                        if len(files) == 0:
                                printf('no file to be operate in ' + root_dir + '\n', 1)
                        for i in files:
                                movefiles(
                                   os.path.normpath(joinp(root_dir, i)),
                                   os.path.normpath(joinp(destdir, root_dir)))
                        if (not KEEP) and len(listdir(root_dir)) == 0 \
                           and root_dir != '.':
                                rmdir(root_dir)
                                printf('clean ' + root_dir + '\n')
        except KeyboardInterrupt:
                exit_err(1, 'User aborted.')

# the main body
#----------- Begin the option parsing -----------#
argv = sys.argv[1:]
opt_hlp, opt_err, opts, params= [], [], [], []
for i in argv:
        if i[0] == '-':
                opts.extend(i[1:])
        else:
                params.append(i)

opt_hlp.append('-n:recursion level(default is 100)\n')
try:
        RECURSION_LEVEL = int(opts[opts.index('-n') + 1])
        del argv[opts.index('-n') + 1]
        del argv[opts.index('-n')]
except:
        RECURSION_LEVEL = 100

opt_hlp.append('-k:keep the empty directories.\n')
if 'k' in opts:
        KEEP = True
else:
        KEEP = False

opt_hlp.append('-m:move the file instead of copying the file\n')
if 'm' in opts:
        OPER_FILE_MOD = 'move'
        OPERATE_FILE = move
else:
        OPER_FILE_MOD = 'copy'
        OPERATE_FILE = copy2

opt_hlp.append('-v:make the output being verbose\n')
if 'v' in opts:
        PRINT_LV = 1
else:
        PRINT_LV = 0

opt_hlp.append('-c:colorize output\n')
if 'c' in opts:
        colours = {'default':"\033[0m",
                   'red'    :"\033[31m",
                   'green'  :"\033[32m",
                   'yellow' :"\033[33m",
                   'blue'   :"\033[34m",
                   }
else:
        colours = {}

try:
        ftypes = params.pop(-1)
        FILE_FILTER = re.compile(ftypes, re.IGNORECASE)
except IndexError:
        opt_err.append('<pattern> missing')
except re.error, er:
        opt_err.append('invalid pattern: ' + ftypes)
        opt_err.append('cause: ' + str(er))

try:
        destdir = params.pop(-1)
except:
        opt_err.append('destdir missing')

try:
        sourcedir = params.pop(-1)
except IndexError:
        sourcedir = '.'

if params:
        opt_err.append('Unknow option: ' + ', '.join(params))

if opt_err != []:
        printf('Usage:sep.py [options] [sourcedir] <destdir> <pattern>\n\
\n\
Find files in sourcedir(default is .) according to pattern and copy/move\n\
them to destdir. destdir must goes before any options. pattern is python\'s\n\
regular expression. Only the last pattern would take effect.\n\
\n\
Options:\n' + '\t' + '\t'.join(opt_hlp)
+ get_color('red') + ', '.join(opt_err) + get_color('default') + '\n')
        sys.exit(2)
#----------- End the option parsing -----------#

printf('ftypes: ' + ftypes + '\n', 1)
printf('destdir: ' + destdir + '\n', 1)

main(sourcedir, destdir)