q-tools/anduril/exec_copy.py

#!/usr/bin/python

import os
import sys
import re
import csv
import shutil
from argparse import ArgumentParser

filesearch=re.compile('^_index$')
STATEFILE='_state'
DRYSTATEFILE='_stateDryRun'

def setup_options():
    usage='''%(prog)s [options]
    Use rsync to copy your execution folder to a new place. 
    e.g.  rsync -avP -e ssh user@server:/source/path/ target/path/
    Then give a replacement pair -i and -o for each expected changed absolute path.
    (multiple -i and -o are often required)
    
    Each replace pair is tried, until a file with that replaced name is found.
    '''

    parser=ArgumentParser(usage=usage)
    parser.add_argument("-d",type=str,dest="execpath",default=".",
                     help="Execution folder for anduril (location for _state), default: %(default)s")
    parser.add_argument("-i",type=str,dest="inabs",action="append",default=["."],
                     help="Input absolute path prefix. e.g. /home1/user1/project/")
    parser.add_argument("-o",type=str,dest="outabs",action="append",default=["."],
                     help="Output absolute path prefix. e.g. /home2/user2/different_project/")
    parser.add_argument("-a",action="store_false",dest="arrays",default=True,
                     help="Disable finding _index (array) files. Just modify _state and nothing else." )
    return parser.parse_args()
    
def check_options(opts):
    if not (os.path.isfile(os.path.join(opts.execpath,STATEFILE))):
        print(STATEFILE+' file not found in folder '+opts.execpath)
        sys.exit(1)
    if len(opts.inabs) is not len(opts.outabs):
        print('A matching pair must be found for each -i/-o argument')
        sys.exit(1)
    if not opts.execpath.endswith('/'):
         opts.execpath=opts.execpath+'/'
    for i in xrange(len(opts.inabs)):
        if not opts.inabs[i].endswith('/'):
            opts.inabs[i]=opts.inabs[i]+'/'
    for i in xrange(len(opts.outabs)):
        if not opts.outabs[i].endswith('/'):
            opts.outabs[i]=opts.outabs[i]+'/'
           
    return opts
    
def getpathlist(path):
    ''' Returns a list of subfolders '''
    list=os.listdir(path)
    paths=[]
    for d in list:
        if (os.path.isdir(os.path.join(path,d))):
            paths.append(d+'/')
    return paths  

def getfilelist(path):
    ''' Returns a list of files that might require change '''
    list=os.listdir(path)
    files=[]
    for f in list:
        if (filesearch.match(f)) and (os.path.isfile(os.path.join(path,f))):
            files.append(f)
    return files    

def statefile(opts,path):
    print('Parsing _state file')
    shutil.copy2(os.path.join(path,STATEFILE), os.path.join(path,STATEFILE+'.bkp'))
    statereader=csv.reader(open(os.path.join(path,STATEFILE),'rb'),
                        delimiter='\t',
                        doublequote=False,
                        escapechar='\\',
                        quoting=csv.QUOTE_NONE)
    stateout=[]
    for row in statereader:
        rowout=row
        if row[3].startswith('INPUT '):
            newinput=row[3]
            rowpstart=row[3].index(' P path=')
            rowtstart=row[3].index(' TS in=')
            rowpath=row[3][(rowpstart+7):rowtstart]
            rowtime=row[3][(rowtstart+7):-1]
            # time has a space at the end (and three zeros...)
            print('INPUT found: "'+row[3]+'"')
            found=False
            for i in xrange(len(opts.inabs)):
                newpath=rowpath.replace('='+opts.inabs[i],'='+opts.outabs[i])
                if os.path.exists(newpath[1:]):
                    found=True
                    newtime=str(int(os.path.getmtime(newpath[1:])))+'000'
                    newinput=row[3].replace(rowpath,newpath,1).replace(rowtime,newtime,1)
                    print('NEW INPUT  : "'+newinput+'"')
                    rowout[3]=newinput
                    break
            if not found:
                print('WARN: Could not find new INPUT, check your -i and -o arguments')
        stateout.append(rowout)    
    statewriter=csv.writer(open(os.path.join(path,STATEFILE),'wb'),
                        delimiter='\t',
                        doublequote=False,
                        escapechar='\\',
                        quoting=csv.QUOTE_NONE)
    statewriter.writerows(stateout)
        
    return

def arrayreplace(path,filelist,opts):
    header=['Key','File']
    for f in filelist:
        print('Modifying array: '+os.path.join(path,f))
        arrayreader=csv.DictReader(open(os.path.join(path,f),'rb'),
                               delimiter='\t',
                               quotechar='"', 
                               quoting=csv.QUOTE_ALL)
        arrayout=[]
        for row in arrayreader:
            rowout=row
            if os.path.exists(os.path.join(path,row['File'])):
                arrayout.append(rowout)
                continue
                # File is a relative path, and exists - next iteration.
            rowpath=row['File']
            found=False
            for i in xrange(len(opts.inabs)):
                newpath=rowpath.replace(opts.inabs[i],opts.outabs[i],1)
                if os.path.exists(newpath):
                    found=True
                    rowout['File']=newpath
                    break
            if not found:
                print('WARN: Could not find File '+rowpath+' in '+os.path.join(path,f)+', check your -i and -o arguments')
            arrayout.append(rowout)
        
        writer = csv.DictWriter(open(os.path.join(path,f),'wb'),
                            header,
                            delimiter='\t',
                            quotechar='"',
                            quoting=csv.QUOTE_MINIMAL)
        writer.writerow(dict(zip(header,header)))
        writer.writerows(arrayout)
        
    return
    
def statequery(path):
    statereader=csv.reader(open(os.path.join(path,DRYSTATEFILE),'rb'),
                        delimiter='\t',
                        doublequote=False,
                        escapechar='\\',
                        quoting=csv.QUOTE_NONE)
    for row in statereader:
        if row[1]=='NO':
            print('Instance will run: '+row[0])
        else:
            print('Instance wont run: '+row[0])
   
    return
    
        
    
def traverse(opts,path):
    pathlist=getpathlist(path)
    filelist=getfilelist(path)
    arrayreplace(path,filelist,opts)
    for p in pathlist:
        traverse(opts,os.path.join(path,p))
    return
    
def main():
    opts=setup_options()
    opts=check_options(opts)

    if opts.arrays:
        traverse(opts,opts.execpath)
    statefile(opts,opts.execpath)
    return


main()