makeTreeClassBase.py

#!/usr/bin/env python
#______________________________________________________________
# makeTreeClassBase.py
# 
# This should be called within the macros/ directory, giving
# a rootfile containing the desired version of the Ntuples as
# an argument.
# 
# A second optional argument can be given to reduce the number
# of branches. This argument is the name of a text file containing
# the desired branches (one branch per line), as produced by the
# getBranches.pl script. Only branches that are in this list are
# then added to the output (only the branches that are already
# in the rootfile are added).
# 
# The outputs are the TreeClassBase header and implementation files.
#______________________________________________________________


import sys, subprocess, os,re
from datetime import datetime
#from ROOT import TTree, TFile, gDirectory

usage = "Usage: makeTreeClassBase.py filename.root [list]"

if len(sys.argv) < 2:
    print usage
    exit(1)

FILENAME = sys.argv[1]
INCLIST = ''
if len(sys.argv)>2: INCLIST = sys.argv[2]
CLASSNAME = 'TreeClassBase'
HEADERNAME = CLASSNAME + '.hh'
SOURCENAME = CLASSNAME + '.cc'

#______________________________________________________________
def isVector(typename):
    return (typename[-1] == 's')

#______________________________________________________________
def typename(vartype):
    typename = vartype.lower() # protect against "Strings"
    # Special case for strings
    if typename.find('string')>-1:
        typename = 'std::string'

    if isVector(vartype): return 'std::vector<'+typename.rstrip('s')+'>'
    else: return typename
        
#______________________________________________________________
def declareVars(names,file):
    indent = 35 # Formatting attempt
    for k,v in sorted(names.iteritems()):
        type = typename(v)
	name = k[0] if k[0] is not "" else k[1]
        # Variable
        spaces = (indent-len(type))*' '+' '
        file.write(4*' '+type+spaces+' '+name+';\n')
        # edm::Handle
        spaces = (indent-14-len(type))*' '+' '
        file.write(4*' '+'edm::Handle<'+type)
        if isVector(v): file.write(' >'+spaces+'h'+name+';\n')
        else: file.write('> '+spaces+'h'+name+';\n')
        # edm::InputTag
        spaces = (indent-13)*' '+' '
        file.write(4*' '+'edm::InputTag'+spaces+'t'+name+';\n')

#______________________________________________________________
def getVars(names,file,spaces,treename):
    indent = spaces*' '
    for k,v in sorted(names.iteritems()):
	name = k[0] if k[0] is not "" else k[1]
        line = 'result &= '+treename+'getByLabel( t'+name+', h'+name+' );\n'
        file.write(indent+line)
        line = 'if ( h'+name+'.isValid() ) '+name+' = *h'+name+';\n'
        file.write(indent+line)
        

#______________________________________________________________
def defineLabels(names,file):
    indent = 25
    for k,v in sorted(names.iteritems()):
	name = k[0] if k[0] is not "" else k[1]
        spaces = (indent-len(name))*' '
        file.write(4*' '+'t'+name+spaces+' = edm::InputTag("'+k[1]+'","'+k[0]+'");\n')

#______________________________________________________________
def processImpl(rBranches,eBranches):
    finput = open('src/base/'+SOURCENAME+'.tpl')
    foutput = open(SOURCENAME,'w')
    patGet    = ('\s*<GET(\w+)HANDLES>.*')
    patLabels = ('\s*<DEFINELABELS>.*')

    today = datetime.today().ctime()
    foutput.write('// This file was automatically generated by '+os.getlogin())
    foutput.write('\n// '+today+'\n')
    for line in finput.readlines():
        m1 = re.match(patGet,line)
        m2 = re.match(patLabels,line)
        if m1:
            if m1.group(1) == 'RUN':
                getVars(rBranches,foutput,6,'run.')
            elif m1.group(1) == 'EVENT':
                getVars(eBranches,foutput,4,'event->')
            else:
                print >>sys.stderr,'*** Unknown header pattern:',line
        elif m2:
            defineLabels(rBranches,foutput)
            defineLabels(eBranches,foutput)
        else:
            foutput.write(line)

    finput.close()
    foutput.close()
    return foutput.name

#______________________________________________________________
def processHeader(rBranches,eBranches):
    finput = open('include/base/'+HEADERNAME+'.tpl')
    foutput = open(HEADERNAME,'w')
    pattern = ('\s*<(\w+)LEAFDECLARATION>.*')

    today = datetime.today().ctime()
    foutput.write('// This file was automatically generated by '+os.getlogin())
    foutput.write('\n// '+today+'\n')
    for line in finput.readlines():
        m = re.match(pattern,line)
        if m:
            if m.group(1) == 'RUN':
                declareVars(rBranches,foutput)
            elif m.group(1) == 'EVENT':
                declareVars(eBranches,foutput)
            else:
                print >>sys.stderr,'*** Unknown header pattern:',line
        else:
            foutput.write(line)

    finput.close()
    foutput.close()
    return foutput.name
            
    
    
#______________________________________________________________
def getBranches(file,tree,include=''):
    cmd = ['edmFileUtil','-P','-t',tree,file]
    run = subprocess.Popen(cmd,stdout=subprocess.PIPE)
    output = run.communicate()[0]
    if run.returncode:
        print >>sys.stderr,"*** Error while parsing file:",output
        return []

    # If list of branches to include is provided, check against it
    docheck = (len(include)>0)
    if docheck:
        list = open(include).read().splitlines()

    branches = dict()
    pattern  = re.compile(r".*?(\w+)_(\S+)_(\S*)_NTupleProducer.*")
    for line in output.split('\n'):
        m  = re.match(pattern,line)
        if m:
            type = m.group(1)
            module = m.group(2)
            branch = m.group(3)
            if docheck and (list.count(branch)==0 and list.count(module)==0): continue
            branches[(branch,module)] = type

	
    return branches

#______________________________________________________________
if __name__=='__main__':

    print 'Processing input file...'
    runBranches = getBranches(FILENAME,'Runs')
    eventBranches = getBranches(FILENAME,'Events',INCLIST)
    print '  -> Found',len(runBranches),'run branches',
    print 'and',len(eventBranches),'event branches'

    if not (len(runBranches)>0 and len(eventBranches)>0):
        print '  *** Stopping here'
        sys.exit(-1)

    # Process templates and write to output files
    hName = processHeader(runBranches,eventBranches)
    print '  -> Wrote',hName
    iName = processImpl(runBranches,eventBranches)
    print '  -> Wrote',iName
	
    subprocess.call(['mv', '-vf', hName, 'include/base/'])
    subprocess.call(['mv', '-vf', iName, 'src/base/'])