| Revision: | 1.19 |
| Committed: | Fri Apr 3 21:00:46 2009 UTC (16 years, 1 month ago) by hegner |
| Content type: | text/x-python |
| Branch: | MAIN |
| CVS Tags: | CMSSW_3_8_6_patch1, CMSSW_3_8_6, CMSSW_3_8_5_patch3, CMSSW_3_8_5_patch2, CMSSW_3_8_4_patch4, CMSSW_3_8_5, CMSSW_3_8_4_patch3, CMSSW_3_8_4_patch2, CMSSW_3_8_4_patch1, CMSSW_3_8_4, CMSSW_3_8_3, CMSSW_3_8_2_patch1, CMSSW_3_6_3_SLHC1, CMSSW_3_8_1_patch4, CMSSW_3_8_2, CMSSW_3_8_1_patch3, CMSSW_3_6_1_patch7, CMSSW_3_8_1_patch2, CMSSW_3_8_1_patch1, CMSSW_3_7_1, CMSSW_3_8_1, CMSSW_3_6_3_patch2, CMSSW_3_8_0_patch1, CMSSW_3_8_0, CMSSW_3_6_1_patch6, CMSSW_3_8_0_pre8, CMSSW_3_8_0_pre7, CMSSW_3_8_0_pre6, CMSSW_3_6_3_patch1, CMSSW_3_7_0_patch4, CMSSW_3_7_0_patch3, CMSSW_3_8_0_pre5, CMSSW_3_6_1_patch5, CMSSW_3_8_0_pre4, CMSSW_3_6_3, CMSSW_3_8_0_pre2, CMSSW_3_7_0_patch2, CMSSW_3_6_2, CMSSW_3_7_0_patch1, CMSSW_3_8_0_pre1, CMSSW_3_6_1_patch4, CMSSW_3_7_0, CMSSW_3_6_1_patch3, CMSSW_3_5_7_hltpatch4, CMSSW_3_7_0_pre5, CMSSW_3_6_1_patch1, CMSSW_3_5_8_patch4, CMSSW_3_7_0_pre4, CMSSW_3_6_1, CMSSW_3_7_0_pre3, CMSSW_3_6_0_patch2, CMSSW_3_6_0_patch1, CMSSW_3_5_8_patch3, CMSSW_3_5_8_patch2, CMSSW_3_5_8_patch1, CMSSW_3_7_0_pre2, CMSSW_3_5_8, CMSSW_3_7_0_pre1, CMSSW_3_5_7, CMSSW_3_6_0, CMSSW_3_6_0_pre6, CMSSW_3_6_0_pre5, CMSSW_3_5_6_patch1, CMSSW_3_5_6, CMSSW_3_5_4_patch2, CMSSW_3_5_5, CMSSW_3_6_0_pre4, CMSSW_3_5_4_patch1, CMSSW_3_6_0_pre3, CMSSW_3_5_4, CMSSW_3_5_3, CMSSW_3_6_0_pre2, CMSSW_3_5_2_patch2, CMSSW_3_5_2_patch1, CMSSW_3_6_0_pre1, CMSSW_3_5_2, CMSSW_3_5_1_patch1, CMSSW_3_3_6_patch4, CMSSW_3_3_6_patch6, CMSSW_3_5_0_patch1, CMSSW_3_4_2_patch1, CMSSW_3_5_1, CMSSW_3_4_2, CMSSW_3_5_0, CMSSW_3_3_6_patch5, CMSSW_3_5_0_pre5, CMSSW_3_3_6_patch3, CMSSW_3_3_3_TSG, CMSSW_3_5_0_pre3, V00-03-01, V00-03-00, CMSSW_3_5_0_pre2, CMSSW_3_4_1, CMSSW_3_4_0, CMSSW_3_3_6_patch2, CMSSW_3_3_6_patch1, CMSSW_3_5_0_pre1, CMSSW_3_3_6, CMSSW_3_4_0_pre7, CMSSW_3_3_5, CMSSW_3_4_0_pre6, CMSSW_3_3_4, CMSSW_3_3_3_patch1, CMSSW_3_3_3, CMSSW_3_4_0_pre5, CMSSW_3_4_0_pre4, CMSSW_3_3_2, CMSSW_3_4_0_pre3, CMSSW_3_3_1, CMSSW_3_4_0_pre2, CMSSW_3_2_8, CMSSW_3_3_0, CMSSW_3_4_0_pre1, CMSSW_3_1_4, CMSSW_3_3_0_pre6, CMSSW_3_3_0_pre5, CMSSW_3_1_3, CMSSW_3_2_7, CMSSW_3_3_0_pre4, CMSSW_3_3_0_pre3, CMSSW_3_3_0_pre2, CMSSW_3_2_6, CMSSW_3_3_0_pre1, CMSSW_3_2_5, CMSSW_3_2_4, CMSSW_3_2_3, V00-02-06, CMSSW_3_2_2_patch2, CMSSW_3_2_2_patch1, CMSSW_3_2_2, CMSSW_3_1_2, CMSSW_3_2_1, CMSSW_3_2_0, CMSSW_3_1_1, CMSSW_2_2_13_offpatch1, CMSSW_3_1_0_patch1, CMSSW_3_1_0, CMSSW_3_1_0_pre11, CMSSW_3_1_0_pre10, V00-02-05, CMSSW_2_2_13_HLT, CMSSW_2_2_13, CMSSW_2_2_12_HLT, CMSSW_2_2_12, CMSSW_3_1_0_pre9, CMSSW_3_1_0_pre8, CMSSW_2_2_11_offpatch1, CMSSW_2_2_11, CMSSW_2_2_11_HLT, CMSSW_3_1_0_pre7, CMSSW_2_2_10_HLT, CMSSW_2_2_10, CMSSW_2_2_9, CMSSW_3_1_0_pre6, CMSSW_2_2_8, CMSSW_3_1_0_pre5, V00-02-04 |
| Changes since 1.18: | +1 -1 lines |
| Log Message: | force ROOT to use batch mode |
| # | 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] != '.') and (branch.GetName()!="EventAuxiliary"): |
| 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 SetAlias (self, alias, fullName): |
| 106 | self.tree().SetAlias(alias, fullName) |
| 107 | def index(self): |
| 108 | return self._index |
| 109 | def tree(self): |
| 110 | return self._tree |
| 111 | def __setBranchIndicies(self): |
| 112 | for branch in self._usedBranches.itervalues(): |
| 113 | branch.setIndex(self._index) |
| 114 | def __getattr__(self, name): |
| 115 | return self.branch(name) |
| 116 | def __getitem__(self,key): |
| 117 | if key <0 or key > self._tree.GetEntries(): |
| 118 | raise IndexError |
| 119 | self._index = key |
| 120 | self.__setBranchIndicies() |
| 121 | self._tree.GetEntry(self._index,0) |
| 122 | return Event(self) |
| 123 | def __iter__(self): |
| 124 | # flushing/initializing the root buffers |
| 125 | entry = 0 |
| 126 | self._index = entry |
| 127 | self.__setBranchIndicies() |
| 128 | self._tree.GetEntry(self._index,0) |
| 129 | # the real loop |
| 130 | for entry in xrange(self._tree.GetEntries()): |
| 131 | self._index = entry |
| 132 | self.__setBranchIndicies() |
| 133 | self._tree.GetEntry(self._index,0) |
| 134 | yield Event(self) # TODO: don't return a new object but update the old one |
| 135 | |
| 136 | |
| 137 | class Event(object): |
| 138 | def __init__(self, eventTree): |
| 139 | self._eventTree = eventTree |
| 140 | |
| 141 | def getProduct(self, name): |
| 142 | return iterators.addIterator(self._eventTree.branch(name)()) |
| 143 | |
| 144 | def __getattr__(self, name): |
| 145 | return iterators.addIterator(self._eventTree.branch(name)()) |
| 146 | |
| 147 | |
| 148 | class EventBranch(object): |
| 149 | def __init__(self,parent,name): |
| 150 | self._branch = parent.tree().GetBranch(name) |
| 151 | if self._branch == None: |
| 152 | raise cmserror("Unknown branch "+name) |
| 153 | self._buffer = createBranchBuffer(self._branch) |
| 154 | self._index = parent.index() |
| 155 | self._readData = False |
| 156 | def setIndex(self,index): |
| 157 | self._index = index |
| 158 | self._readData = False |
| 159 | def __readData(self): |
| 160 | self._branch.GetEntry(self._index) |
| 161 | self._readData = True |
| 162 | |
| 163 | # replace this by __getattr__ to allow branch.attr instead of branch().attr |
| 164 | def __call__(self): |
| 165 | if not self._readData: |
| 166 | self.__readData() |
| 167 | return self._buffer |
| 168 | |
| 169 | |
| 170 | class cmserror(exceptions.StandardError): |
| 171 | def __init__(self, message): |
| 172 | length = len(message)+7 #7=len("ERROR: ") |
| 173 | print "="*length |
| 174 | print "ERROR:", message |
| 175 | print "="*length |