Ejemplo n.º 1
0
def gardener_cli( modules ):
    usage = '''
    Usage:
        %prog <command> <options> filein.root fileout.root
        %prog <command> <options> file1.root file2.root ... dirout
        %prog <command> -r <options> dirin dirout

    In the latter case the directory tree in dirin is rebuilt in dirout

    Valid commands:
        '''+', '.join(modules.keys()+['help'])+'''

    Type %prog <command> -h for the command specific help
    '''

    parser = optparse.OptionParser(usage)
    parser.add_option('-t','--tree',        dest='tree',                                default='latino',   help='Name of the tree to operate on (default = %default)')
    parser.add_option('-r','--recursive',   dest='recursive',   action='store_true',    default=False,      help='Recurse subdirectories (default = %default)')
    parser.add_option('-F','--force',       dest='force',       action='store_true',    default=False,      help='Don\'t ask for confirmation when recursing (default = %default)')

    # some boring argument handling
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)

    modname = sys.argv[1]
    if modname.startswith('-'):
        parser.print_help()
        sys.exit(0)

    if modname == 'help':
        if len(sys.argv) == 2:
            parser.print_help()
            sys.exit(0)
        
        module = sys.argv[2]
        if module not in modules:
            print 'Help: command',module,'not known'
            print 'The available commands are',', '.join(modules.keys())
            sys.exit(0)
        
        print 'Help for module',module+':'
        modules[module].help()
        modules[module].addOptions(parser)
        parser.print_help()
        sys.exit(0)


    if modname not in modules:
        print 'Command',modname,'unknown'
        print 'The available commands are',modules.keys()
        sys.exit(0)

    module = modules[modname]
    group = module.addOptions(parser)

    sys.argv.remove(modname)

    (opt, args) = parser.parse_args()

    print opt,args

    sys.argv.append('-b')

    try:
        module.checkOptions(opt)
    except Exception as e:
        print 'Error in module',module.label
#         print '*'*80
#         print 'Fatal exception '+type(e).__name__+': '+str(e)
#         print '*'*80
#         exc_type, exc_value, exc_traceback = sys.exc_info()
#         traceback.print_tb(exc_traceback, file=sys.stdout)
#         print '*'*80
        print e
        sys.exit(1)

    tree = opt.tree

    nargs = len(args)
    if nargs < 2:
        parser.error('Input and/or output files missing')


    # file1 file2 file3 > outdir
    elif nargs > 2:
        inputs = args[:-1]
        output = args[-1]

        print inputs
        
        # sanitise output
        output = output if output[-1]=='/' else output+'/'
        iofiles = [ (f,os.path.join(output,os.path.basename(f))) for f in inputs ]

        execute( module, tree, iofiles )
    
    elif nargs == 2:
        input  = args[0]
        output = args[1]


        # recursiveness here
        if os.path.isdir(input):
            if not opt.recursive:
                print input,'is a directory. Use -r to go recursive'
                sys.exit(0)

            # sanitize the input/output
            input  = input  if input [-1]=='/' else input +'/'
            output = output if output[-1]=='/' else output+'/'

            if os.path.exists(output) and not os.path.isdir(output):
                print output,'exists and is not a directory!'
                sys.exit(0)

            fileList = []
            for root, subFolders, files in os.walk(input):
                for file in files:
                    if not file.endswith('.root'): continue
                    fileList.append(os.path.join(root,file))

            print 'The directory tree',input,'will be gardened and copied to',output
            print 'The following files will be copied:'
            print '\n'.join(fileList)
            print 'for a grand total of',len(fileList),'files'
            opt.force or ( confirm('Do you want to continue?') or sys.exit(0) )

            iofiles = [ (f,f.replace(input,output)) for f in fileList ]

            execute( module, tree, iofiles )

        else:
            if os.path.exists(output) and os.path.isdir(output):
                # sanitise output
                output = output if output[-1]=='/' else output+'/'
                output = os.path.join( output , os.path.basename(input) )
            execute(module,tree,[(input,output)])
Ejemplo n.º 2
0
    def process(self,**kwargs):

        tree  = kwargs['tree']
        input = kwargs['input']
        output = kwargs['output']

        self.connect(tree,input)

        vars = [ ( value, type, ROOT.TTreeFormula(name,formula, self.itree)) for name, (value, type, formula) in self.variables.iteritems() ]


        print 'Adding/replacing the following branches'
        template=' {0:10} | {1:^3} | "{2}"'
        for name  in sorted(self.variables):
            (value, type, formula) = self.variables[name]
            print template.format(name,type,formula)
        print


        oldbranches = [ b.GetName() for b in self.itree.GetListOfBranches() ]
        hasMutation = False
        for bname in self.variables:
            # not there, continue
            if bname not in oldbranches: continue
            # found, check for consistency
            branch = self.itree.GetBranch(bname)
            btype = self.variables[bname][1]
            newtitle = bname+'/'+btype
            if ( branch.GetTitle() != newtitle ):
                print 'WARNING: Branch mutation detected: from',branch.GetTitle(),'to',newtitle
                hasMutation = True

        if hasMutation:
            confirm('Mutation detected. Do you _really_ want to continue?') or sys.exit(0)

        self.clone(output,self.variables.keys())

        for (val,type,formula) in vars:
            name = formula.GetName()
            title = name+'/'+type
            self.otree.Branch(name,val,title)

        nentries = self.itree.GetEntries()
        print 'Entries:',nentries

        # avoid dots in the loop
        itree = self.itree
        otree = self.otree

        step = 5000
        for i in xrange(0,nentries):
            itree.GetEntry(i)

            if i > 0 and i%step == 0:
                print str(i)+' events processed.'

            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                for (val,type,formula) in vars:
                    val[0] = formula.EvalInstance()

            otree.Fill()
        
        self.disconnect()
        print '- Eventloop completed'