ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/CMSSW/PhysicsTools/PythonAnalysis/python/cmstools.py
Revision: 1.16
Committed: Tue Sep 16 16:15:54 2008 UTC (16 years, 7 months ago) by hegner
Content type: text/x-python
Branch: MAIN
CVS Tags: CMSSW_3_1_0_pre3, CMSSW_3_1_0_pre2, CMSSW_3_0_0_pre10, CMSSW_2_2_5, CMSSW_3_0_0_pre9, CMSSW_3_1_0_pre1, CMSSW_2_2_4, CMSSW_3_0_0_pre8, CMSSW_3_0_0_pre7, CMSSW_3_0_0_pre6, CMSSW_3_0_0_pre5, CMSSW_3_0_0_pre4, CMSSW_2_2_3, CMSSW_3_0_0_pre3, CMSSW_2_2_2, CMSSW_2_2_1, CMSSW_2_2_0, CMSSW_3_0_0_pre2, CMSSW_2_2_0_pre1, CMSSW_2_1_12, CMSSW_3_0_0_pre1, CMSSW_2_1_11, CMSSW_2_1_10_patch2, CMSSW_2_1_10_patch1, CMSSW_2_1_10, V00-02-01
Changes since 1.15: +2 -1 lines
Log Message:
make EventTree.__getitem__ more robust

File Contents

# Content
1 """Python helper tools for CMS FWLite
2
3 benedikt.hegner@cern.ch
4
5 """
6 import re
7 import ROOT
8 import exceptions
9 ### define tab completion
10 try:
11 import readline, cmscompleter
12 readline.parse_and_bind('tab: complete')
13 except:
14 print 'WARNING: Could not load tab completion'
15
16
17 # for adding iterators at runtime
18 import iterators
19
20
21 ### workaround iterator generators for ROOT classes
22 def all(container):
23
24 # loop over ROOT::TTree and similar
25 if hasattr(container,'GetEntries'):
26 try:
27 entries = container.GetEntries()
28 for entry in xrange(entries):
29 yield entry
30 except:
31 raise cmserror("Looping of %s failed" %container)
32
33 # loop over std::vectors and similar
34 elif hasattr(container, 'size'):
35 try:
36 entries = container.size()
37 for entry in xrange(entries):
38 yield container[entry]
39 except:
40 pass
41
42 # loop over containers with begin and end iterators
43 def loop(begin, end):
44 """Convert a pair of C++ iterators into a python generator"""
45 while (begin != end):
46 yield begin.__deref__() #*b
47 begin.__preinc__() #++b
48
49 ### auto branch types (Chris Jones)
50 def createBranchBuffer(branch):
51 reColons = re.compile(r'::')
52 reCloseTemplate =re.compile(r'>')
53 reOpenTemplate =re.compile(r'<')
54 branchType = ROOT.branchToClass(branch)
55 #buffer = eval ('ROOT.'+reColons.sub(".",reOpenTemplate.sub("(ROOT.",reCloseTemplate.sub(")",branchType.GetName())))+'()')
56 buffer = ROOT.MakeRootClass(branchType.GetName()) ()
57 if( branch.GetName()[-1] != '.'):
58 branch.SetAddress(buffer)
59 else:
60 branch.SetAddress(ROOT.AddressOf(buffer))
61 return buffer
62
63
64 class EventTree(object):
65 def __init__(self,obj):
66 if isinstance(obj, ROOT.TTree):
67 self._tree = obj
68 elif isinstance(obj, ROOT.TFile):
69 self._tree = obj.Get("Events")
70 elif isinstance(obj, str):
71 self._tree = ROOT.TFile.Open(obj).Get("Events")
72 else:
73 raise cmserror("EventTree accepts only TTrees, TFiles and filenames")
74 self._usedBranches = dict()
75 self._index = -1
76 self._aliases = self._tree.GetListOfAliases()
77 def branch(self,name):
78 # support for aliases
79 alias = self._tree.GetAlias(name)
80 if alias != '': name = alias
81 # access the branch in ttree
82 if name in self._usedBranches:
83 return self._usedBranches[name]
84 self._usedBranches[name]=EventBranch(self,name)
85 return self._usedBranches[name]
86 def cppCode(self, name):
87 """C++ code for accessing the product inside the full framework"""
88 alias = self._tree.GetAlias(name)
89 if alias != '': name = alias
90 tmpBranch = self._tree.GetBranch(name)
91 typeString = ROOT.branchToClass(tmpBranch).GetName()
92 if "edm::Wrapper" in typeString:
93 typeString = typeString.replace("<edm::Wrapper","")
94 typeString = typeString.rstrip(">")
95 nameParts = name.split("_")
96 if nameParts[2] == "":
97 cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", dummy);'\
98 %(typeString, nameParts[1])
99 else:
100 cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", "%s", dummy);'\
101 %(typeString, nameParts[1], nameParts[2])
102 return cppCode
103 def getListOfAliases(self):
104 return self._aliases
105 def index(self):
106 return self._index
107 def tree(self):
108 return self._tree
109 def __setBranchIndicies(self):
110 for branch in self._usedBranches.itervalues():
111 branch.setIndex(self._index)
112 def __getattr__(self, name):
113 return self.branch(name)
114 def __getitem__(self,key):
115 if key <0 or key > self._tree.GetEntries():
116 raise IndexError
117 self._index = key
118 self.__setBranchIndicies()
119 self._tree.GetEntry(self._index,0)
120 return Event(self)
121 def __iter__(self):
122 # flushing/initializing the root buffers
123 entry = 0
124 self._index = entry
125 self.__setBranchIndicies()
126 self._tree.GetEntry(self._index,0)
127 # the real loop
128 for entry in xrange(self._tree.GetEntries()):
129 self._index = entry
130 self.__setBranchIndicies()
131 self._tree.GetEntry(self._index,0)
132 yield Event(self) # TODO: don't return a new object but update the old one
133
134
135 class Event(object):
136 def __init__(self, eventTree):
137 self._eventTree = eventTree
138
139 def getProduct(self, name):
140 return iterators.addIterator(self._eventTree.branch(name)())
141
142 def __getattr__(self, name):
143 return iterators.addIterator(self._eventTree.branch(name)())
144
145
146 class EventBranch(object):
147 def __init__(self,parent,name):
148 self._branch = parent.tree().GetBranch(name)
149 if self._branch == None:
150 raise cmserror("Unknown branch "+name)
151 self._buffer = createBranchBuffer(self._branch)
152 self._index = parent.index()
153 self._readData = False
154 def setIndex(self,index):
155 self._index = index
156 self._readData = False
157 def __readData(self):
158 self._branch.GetEntry(self._index)
159 self._readData = True
160
161 # replace this by __getattr__ to allow branch.attr instead of branch().attr
162 def __call__(self):
163 if not self._readData:
164 self.__readData()
165 return self._buffer
166
167
168 class cmserror(exceptions.StandardError):
169 def __init__(self, message):
170 length = len(message)+7 #7=len("ERROR: ")
171 print "="*length
172 print "ERROR:", message
173 print "="*length